Merge branch 'main' into mrv2-ci-test

[WIP] CI test run with Model Runner V2
Expected to fail many tests Signed-off-by: Nick Hill <nickhill123@gmail.com>
2026-03-14 10:29:16 -07:00 · 2026-02-27 14:48:27 -08:00
934 changed files with 25429 additions and 48183 deletions
@@ -0,0 +1,12 @@
+# For vllm script, with -t option (tensor parallel size).
+# bash ./run-lm-eval-gsm-vllm-baseline.sh -m nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM -b "auto" -t 2
+model_name: "nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.6353
+  - name: "exact_match,flexible-extract"
+    value: 0.637
+limit: null
+num_fewshot: null 
@@ -1 +0,0 @@
-Qwen3-235B-A22B-Instruct-2507-FP8.yaml
@@ -12,7 +12,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cuda-12-9
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
@@ -27,7 +27,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cuda-13-0
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
          # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
@@ -42,7 +42,7 @@ steps:
        depends_on: ~
        id: build-wheel-arm64-cpu
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
@@ -55,7 +55,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cuda-12-9
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
@@ -68,7 +68,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cuda-13-0
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
@@ -81,7 +81,7 @@ steps:
        depends_on: ~
        id: build-wheel-x86-cpu
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
          - "mkdir artifacts"
@@ -97,7 +97,7 @@ steps:
        depends_on: ~
        id: build-release-image-x86
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -110,7 +110,7 @@ steps:
        depends_on: ~
        id: build-release-image-arm64
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -120,7 +120,7 @@ steps:
        depends_on: ~
        id: build-release-image-x86-cuda-13-0
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
@@ -133,7 +133,7 @@ steps:
        depends_on: ~
        id: build-release-image-arm64-cuda-13-0
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          # compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
@@ -149,7 +149,7 @@ steps:
          - block-cpu-release-image-build
          - input-release-version
        agents:
-          queue: cpu_queue_release
+          queue: cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
@@ -167,7 +167,7 @@ steps:
          - block-arm64-cpu-release-image-build
          - input-release-version
        agents:
-          queue: arm64_cpu_queue_release
+          queue: arm64_cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
@@ -185,7 +185,7 @@ steps:
          - build-release-image-arm64
        id: create-multi-arch-manifest
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
@@ -196,7 +196,7 @@ steps:
          - create-multi-arch-manifest
        id: annotate-release-workflow
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/annotate-release.sh"

@@ -206,7 +206,7 @@ steps:
          - build-release-image-arm64-cuda-13-0
        id: create-multi-arch-manifest-cuda-13-0
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
@@ -217,7 +217,7 @@ steps:
          - create-multi-arch-manifest
        if: build.env("NIGHTLY") == "1"
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/push-nightly-builds.sh"
          # Clean up old nightly builds (keep only last 14)
@@ -235,7 +235,7 @@ steps:
          - create-multi-arch-manifest-cuda-13-0
        if: build.env("NIGHTLY") == "1"
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/push-nightly-builds.sh cu130"
          # Clean up old nightly builds (keep only last 14)
@@ -262,7 +262,7 @@ steps:
          - block-upload-release-wheels
        id: upload-release-wheels
        agents:
-          queue: small_cpu_queue_release
+          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"

@@ -323,7 +323,7 @@ steps:
      - step: input-rocm-config
        allow_failure: true  # Allow failure so non-UI builds can proceed (input step is skipped)
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    commands:
      # Set configuration and check cache
      - |
@@ -465,7 +465,7 @@ steps:
      - step: build-rocm-base-wheels
        allow_failure: false
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    timeout_in_minutes: 180
    commands:
      # Download artifacts and prepare Docker image
@@ -575,7 +575,7 @@ steps:
      - step: build-rocm-vllm-wheel
        allow_failure: false
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    timeout_in_minutes: 60
    commands:
      # Download all wheel artifacts and run upload
@@ -624,7 +624,7 @@ steps:
      - step: input-release-version
        allow_failure: true
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    commands:
      - "bash .buildkite/scripts/annotate-rocm-release.sh"
    env:
@@ -641,7 +641,7 @@ steps:
    depends_on: block-generate-root-index-rocm-wheels
    id: generate-root-index-rocm-wheels
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    commands:
      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
    env:
@@ -655,7 +655,7 @@ steps:
      - step: build-rocm-base-wheels
        allow_failure: false
    agents:
-      queue: cpu_queue_release
+      queue: cpu_queue_postmerge
    timeout_in_minutes: 60
    commands:
      - |
@@ -16,23 +16,6 @@ RAY_BASE_URL="https://raw.githubusercontent.com/ray-project/ray/master/python"
 WORK_DIR=$(mktemp -d)
 trap 'rm -rf "$WORK_DIR"' EXIT

-# ── Detect PyTorch index URL ─────────────────────────────────────────────
-
-if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
-    ROCM_VER=$(python3 -c "import torch; print(torch.version.hip.rsplit('.', 1)[0])")
-    CANDIDATE_URL="https://download.pytorch.org/whl/rocm${ROCM_VER}"
-    if curl -fsSL --head "${CANDIDATE_URL}/" >/dev/null 2>&1; then
-        TORCH_INDEX_URL="${CANDIDATE_URL}"
-    else
-        echo ">>> WARNING: ROCm ${ROCM_VER} wheel index not found at ${CANDIDATE_URL}"
-        echo ">>>          Falling back to default PyPI (resolution may be incomplete)"
-        TORCH_INDEX_URL=""
-    fi
-else
-    TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
-fi
-echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
-
 # Fetch all Ray requirement files used in the LLM depset pipeline
 echo ">>> Fetching Ray requirement files"
 RAY_FILES=(
@@ -133,11 +116,6 @@ echo "============================================================"
 echo ">>> Resolving: Can Ray generate compatible lock files?"
 echo "============================================================"

-EXTRA_INDEX_ARGS=()
-if [[ -n "${TORCH_INDEX_URL}" ]]; then
-    EXTRA_INDEX_ARGS+=(--extra-index-url "${TORCH_INDEX_URL}")
-fi
-
 set +e
 uv pip compile \
    "${WORK_DIR}/requirements.txt" \
@@ -148,7 +126,7 @@ uv pip compile \
    -c "${WORK_DIR}/vllm-constraints.txt" \
    --python-version 3.12 \
    --python-platform x86_64-manylinux_2_31 \
-    "${EXTRA_INDEX_ARGS[@]}" \
+    --extra-index-url https://download.pytorch.org/whl/cu129 \
    --index-strategy unsafe-best-match \
    --unsafe-package setuptools \
    --unsafe-package ray \
@@ -326,24 +326,22 @@ apply_rocm_test_overrides() {
  if [[ $cmds == *" kernels/moe"* ]]; then
    cmds="${cmds} \
    --ignore=kernels/moe/test_moe.py \
-    --ignore=kernels/moe/test_cutlass_moe.py"
+    --ignore=kernels/moe/test_cutlass_moe.py \
+    --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
  fi

  # --- Entrypoint ignores ---
  if [[ $cmds == *" entrypoints/openai "* ]]; then
    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/chat_completion/test_audio.py \
-    --ignore=entrypoints/openai/completion/test_shutdown.py \
+    --ignore=entrypoints/openai/test_audio.py \
+    --ignore=entrypoints/openai/test_shutdown.py \
    --ignore=entrypoints/openai/test_completion.py \
-    --ignore=entrypoints/openai/models/test_models.py \
+    --ignore=entrypoints/openai/test_models.py \
+    --ignore=entrypoints/openai/test_lora_adapters.py \
    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
-    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
-  fi
-
-  if [[ $cmds == *" entrypoints/serve"* ]]; then
-    cmds="${cmds} \
-    --ignore=entrypoints/serve/lora/test_lora_adapters.py"
+    --ignore=entrypoints/openai/test_root_path.py \
+    --ignore=entrypoints/openai/test_tokenization.py \
+    --ignore=entrypoints/openai/test_prompt_validation.py "}
  fi

  if [[ $cmds == *" entrypoints/llm "* ]]; then
@@ -127,7 +127,7 @@ run_and_track_test() {

 # --- Actual Test Execution ---
 run_and_track_test 1 "test_struct_output_generate.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
+    "python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
 run_and_track_test 2 "test_moe_pallas.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
 run_and_track_test 3 "test_lora.py" \
@@ -33,22 +33,23 @@ docker run \
    bash -c '
    set -e
    echo $ZE_AFFINITY_MASK
+    pip install tblib==3.1.0
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
    cd tests
    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
    pytest -v -s v1/engine
    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
-    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
    pytest -v -s v1/structured_output
    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
-    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)"
+    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py
    pytest -v -s v1/test_serial_utils.py
 '
@@ -1,14 +1,11 @@
 #!/usr/bin/env bash
 set -euxo pipefail
+
 # Nightly e2e test for prefetch offloading with a MoE model.
 # Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
 # and validates GSM8K accuracy matches baseline (no offloading).
 #
 # args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
-#
-# Environment variables:
-#   ATTENTION_BACKEND   - attention backend to use (e.g., FLASH_ATTN,
-#                         ROCM_ATTN, FLASHINFER). If unset, uses vllm default.
 THRESHOLD=${1:-0.25}
 NUM_Q=${2:-1319}
 PORT=${3:-8030}
@@ -25,14 +22,6 @@ wait_for_server() {

 MODEL="deepseek-ai/DeepSeek-V2-Lite"

-# ── Build optional vllm serve flags ─────────────────────────────────────
-
-EXTRA_ARGS=()
-if [[ -n "${ATTENTION_BACKEND:-}" ]]; then
-  echo "Using attention backend: ${ATTENTION_BACKEND}"
-  EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
-fi
-
 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
    kill "${SERVER_PID}" 2>/dev/null || true
@@ -51,8 +40,7 @@ vllm serve "$MODEL" \
  --offload-num-in-group 2 \
  --offload-prefetch-step 1 \
  --offload-params w13_weight w2_weight \
-  --port "$PORT" \
-  ${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
+  --port "$PORT" &
 SERVER_PID=$!
 wait_for_server "$PORT"

@@ -59,7 +59,7 @@ steps:
  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
  - pytest -s -v tests/compile/passes/distributed

- label: Fusion and Compile Unit Tests (2xB200)
+- label: Fusion and Compile Unit Tests (B200)
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
  device: b200
@@ -15,29 +15,8 @@ steps:
  - pytest -v -s distributed/test_shm_buffer.py
  - pytest -v -s distributed/test_shm_storage.py

- label: Distributed DP Tests (2 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/v1/distributed
-  - tests/entrypoints/openai/test_multi_api_servers.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
-
- label: Distributed Compile + RPC Tests (2 GPUs)
-  timeout_in_minutes: 20
+- label: Distributed (2 GPUs)
+  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
@@ -50,31 +29,22 @@ steps:
  - vllm/v1/worker/
  - tests/compile/fullgraph/test_basic_correctness.py
  - tests/compile/test_wrapper.py
-  - tests/entrypoints/llm/test_collective_rpc.py
-  commands:
-  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - 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
-
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
  - tests/distributed/
+  - tests/entrypoints/llm/test_collective_rpc.py
+  - tests/v1/distributed
+  - tests/v1/entrypoints/openai/test_multi_api_servers.py
  - tests/v1/shutdown
  - tests/v1/worker/test_worker_memory_snapshot.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -82,35 +52,41 @@ steps:

 - label: Distributed Torchrun + Examples (4 GPUs)
  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace"
+  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
  - vllm/distributed/
  - tests/distributed/test_torchrun_example.py
  - tests/distributed/test_torchrun_example_moe.py
+  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
-  - examples/rl/
+  - examples/offline_inference/new_weight_syncing/
  - tests/examples/offline_inference/data_parallel.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
  # test with torchrun tp=2 and external_dp=2
-  - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  # test with torchrun tp=2 and pp=2
-  - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  # test with torchrun tp=4 and dp=1
-  - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2, pp=2 and dp=1
-  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=1 and dp=4 with ep
-  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2 and dp=2 with ep
-  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
  # test with internal dp
-  - python3 examples/offline_inference/data_parallel.py --enforce-eager
-  # rlhf examples
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
+  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
+  # OLD rlhf examples
+  - cd ../examples/offline_inference
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
+  # NEW rlhf examples
+  - cd new_weight_syncing
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py

 - label: Distributed DP Tests (4 GPUs)
  timeout_in_minutes: 30
@@ -193,7 +169,7 @@ steps:
  num_devices: 2
  commands:
    - pytest -v -s tests/distributed/test_context_parallel.py
-    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
+    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
    - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.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

@@ -257,17 +233,6 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
-  timeout_in_minutes: 20
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - 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
-    - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 - label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
  timeout_in_minutes: 30
  device: a100
@@ -70,15 +70,3 @@ steps:
      device: mi325_4
      depends_on:
      - image-build-amd
-
- label: V1 e2e (4xH100)
-  timeout_in_minutes: 60
-  device: h100
-  num_devices: 4
-  optional: true
-  source_file_dependencies:
-    - vllm/v1/attention/backends/utils.py
-    - vllm/v1/worker/gpu_model_runner.py
-    - tests/v1/e2e/test_hybrid_chunked_prefill.py
-  commands:
-    - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py
@@ -10,7 +10,7 @@ steps:
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration (LLM)
  timeout_in_minutes: 40
@@ -25,8 +25,8 @@ steps:
  - 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

- label: Entrypoints Integration (API Server openai - Part 1)
-  timeout_in_minutes: 50
+- label: Entrypoints Integration (API Server 1)
+  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -34,24 +34,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  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)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai
-  - tests/entrypoints/test_chat_utils
-  commands:
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - pytest -v -s entrypoints/openai/speech_to_text/
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py
  mirror:
    amd:
@@ -59,28 +42,17 @@ steps:
      depends_on:
      - image-build-amd

- label: Entrypoints Integration (API Server openai - Part 3)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai
-  - tests/entrypoints/test_chat_utils
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
-
 - label: Entrypoints Integration (API Server 2)
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/rpc
-  - tests/entrypoints/serve/instrumentator
+  - tests/entrypoints/instrumentator
  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/serve/instrumentator
+  - pytest -v -s entrypoints/instrumentator
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use

@@ -103,6 +75,19 @@ steps:
  commands:
  - pytest -v -s entrypoints/openai/responses

+- label: Entrypoints V1
+  timeout_in_minutes: 50
+  source_file_dependencies:
+    - vllm/
+    - tests/v1
+  commands:
+    - pytest -v -s v1/entrypoints
+  mirror:
+    amd:
+      device: mi325_1
+      depends_on:
+      - image-build-amd
+
 - label: OpenAI API Correctness
  timeout_in_minutes: 30
  source_file_dependencies:
@@ -24,7 +24,8 @@ steps:

 - label: Elastic EP Scaling Test
  timeout_in_minutes: 20
-  device: h100
+  device: b200
+  optional: true
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -35,7 +35,7 @@ steps:
  parallelism: 2

 - label: Kernels MoE Test %N
-  timeout_in_minutes: 25
+  timeout_in_minutes: 60
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -47,7 +47,7 @@ steps:
  commands:
    - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 5
+  parallelism: 2

 - label: Kernels Mamba Test
  timeout_in_minutes: 45
@@ -45,22 +45,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

- label: LM Eval Qwen3.5 Models (B200)
-  timeout_in_minutes: 120
-  device: b200
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - vllm/model_executor/models/qwen3_5.py
-  - vllm/model_executor/models/qwen3_5_mtp.py
-  - vllm/transformers_utils/configs/qwen3_5.py
-  - vllm/transformers_utils/configs/qwen3_5_moe.py
-  - vllm/model_executor/models/qwen3_next.py
-  - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/model_executor/layers/fla/ops/
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
-
 - label: LM Eval Large Models (H200)
  timeout_in_minutes: 60
  device: h200
@@ -90,7 +74,6 @@ steps:
  commands:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt

-
 - label: GPQA Eval (GPT-OSS) (H100)
  timeout_in_minutes: 120
  device: h100
@@ -8,7 +8,7 @@ steps:
  - vllm/lora
  - tests/lora
  commands:
-    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
+    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py
  parallelism: 4


@@ -30,5 +30,4 @@ steps:
    - pytest -v -s -x lora/test_llama_tp.py
    - pytest -v -s -x lora/test_llm_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
-    - pytest -v -s -x lora/test_gptoss_tp.py
-    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+    - pytest -v -s -x lora/test_gptoss_tp.py
@@ -2,54 +2,11 @@ group: Miscellaneous
 depends_on: 
  - image-build
 steps:
- label: V1 Spec Decode
-  timeout_in_minutes: 30
+- label: V1 Others
+  timeout_in_minutes: 60
  source_file_dependencies:
    - 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
-  timeout_in_minutes: 30
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/sample
-    - tests/v1/logits_processors
-    - tests/v1/test_oracle.py
-    - tests/v1/test_request.py
-    - tests/v1/test_outputs.py
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s v1/sample
-    - 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
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-
- label: V1 Core + KV + Metrics
-  timeout_in_minutes: 30
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/core
-    - tests/v1/executor
-    - tests/v1/kv_offload
-    - tests/v1/worker
-    - tests/v1/kv_connector/unit
-    - tests/v1/metrics
-    - tests/entrypoints/openai/correctness/test_lmeval.py
+    - tests/v1
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -57,9 +14,16 @@ steps:
    - pytest -v -s -m 'not cpu_test' v1/core
    - pytest -v -s v1/executor
    - pytest -v -s v1/kv_offload
+    - pytest -v -s v1/sample
+    - pytest -v -s v1/logits_processors
    - pytest -v -s v1/worker
+    # TODO: create another `optional` test group for slow tests
+    - pytest -v -s -m 'not slow_test' v1/spec_decode
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
+    - pytest -v -s v1/test_oracle.py
+    - pytest -v -s v1/test_request.py
+    - pytest -v -s v1/test_outputs.py
    # 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
@@ -75,7 +39,7 @@ steps:
  source_file_dependencies:
    - vllm/
    - tests/v1
-  device: cpu-small
+  device: cpu
  commands:
    # split the test to avoid interference
    - pytest -v -s -m 'cpu_test' v1/core
@@ -177,7 +141,7 @@ steps:
  - tests/tool_parsers
  - tests/transformers_utils
  - tests/config
-  device: cpu-small
+  device: cpu
  commands:
  - python3 standalone_tests/lazy_imports.py
  - pytest -v -s test_inputs.py
@@ -192,7 +156,7 @@ steps:
  - pytest -v -s config

 - label: Batch Invariance (H100)
-  timeout_in_minutes: 30
+  timeout_in_minutes: 25
  device: h100
  source_file_dependencies:
    - vllm/v1/attention
@@ -203,8 +167,6 @@ steps:
    - pip install pytest-timeout pytest-forked
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
  
 - label: Acceptance Length Test (Large Models) # optional
  timeout_in_minutes: 25
@@ -9,9 +9,9 @@ steps:
  - vllm/config/model.py
  - vllm/model_executor
  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - tests/entrypoints/openai/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -v -s model_executor
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+    - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
@@ -11,7 +11,7 @@ steps:
  - vllm/v1/attention/
  - tests/v1/engine/test_llm_engine.py
  - tests/v1/e2e/
-  - tests/entrypoints/llm/test_struct_output_generate.py
+  - tests/v1/entrypoints/llm/test_struct_output_generate.py
  commands:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
@@ -22,7 +22,7 @@ steps:
  - pytest -v -s v1/e2e/general/test_context_length.py
  - pytest -v -s v1/e2e/general/test_min_tokens.py
  # Temporary hack filter to exclude ngram spec decoding based tests.
-  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
+  - pytest -v -s v1/entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"

 - label: Model Runner V2 Examples
  timeout_in_minutes: 45
@@ -87,12 +87,13 @@ steps:
    - vllm/v1/worker/gpu/
    - vllm/v1/worker/gpu_worker.py
    - tests/distributed/test_pipeline_parallel.py
-    - tests/distributed/test_pp_cudagraph.py
+    #- tests/distributed/test_pp_cudagraph.py
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
    - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
-    - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
+    # TODO: Uncomment once https://github.com/vllm-project/vllm/pull/35162 is merged.
+    #- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"

 - label: Model Runner V2 Spec Decode
  timeout_in_minutes: 30
@@ -51,7 +51,7 @@ steps:
  - vllm/
  - tests/models/test_utils.py
  - tests/models/test_vision.py
-  device: cpu-small
+  device: cpu
  commands:
    - pytest -v -s models/test_utils.py models/test_vision.py

@@ -62,7 +62,7 @@ steps:
      depends_on:
      - image-build-amd

- label: Multi-Modal Processor (CPU)
+- label: Multi-Modal Processor Test (CPU)
  depends_on: 
  - image-build-cpu
  timeout_in_minutes: 60
@@ -70,7 +70,7 @@ steps:
  - vllm/
  - tests/models/multimodal
  - tests/models/registry.py
-  device: cpu-medium
+  device: cpu
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
@@ -95,44 +95,34 @@ steps:
  commands:
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1

- label: Multi-Modal Models (Extended Generation 1)
+- label: Multi-Modal Models (Extended) 1
  optional: true
  source_file_dependencies:
  - vllm/
-  - tests/models/multimodal/generation
-  - tests/models/multimodal/test_mapping.py
+  - tests/models/multimodal
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
-    - pytest -v -s models/multimodal/test_mapping.py
+    - pytest -v -s models/multimodal -m 'not core_model' --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing
  mirror:
    amd:
      device: mi325_1
      depends_on:
      - image-build-amd

- label: Multi-Modal Models (Extended Generation 2)
+- label: Multi-Modal Models (Extended) 2
  optional: true
  source_file_dependencies:
  - vllm/
-  - tests/models/multimodal/generation
+  - tests/models/multimodal
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'

- label: Multi-Modal Models (Extended Generation 3)
+- label: Multi-Modal Models (Extended) 3
  optional: true
  source_file_dependencies:
  - vllm/
-  - tests/models/multimodal/generation
+  - tests/models/multimodal
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
-
- label: Multi-Modal Models (Extended Pooling)
-  optional: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal/pooling
-  commands:
-    - pytest -v -s models/multimodal/pooling -m 'not core_model'
@@ -36,6 +36,6 @@ steps:
  - pytest -v -s plugins_tests/test_scheduler_plugins.py
  - pip install -e ./plugins/vllm_add_dummy_model
  - pytest -v -s distributed/test_distributed_oot.py
-  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
+  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
@@ -35,7 +35,7 @@ steps:
  # as it is a heavy test that is covered in other steps.
  # Use `find` to launch multiple instances of pytest so that
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
+  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\;"

 - label: PyTorch Fullgraph
  timeout_in_minutes: 30
@@ -75,7 +75,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
 /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/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
 /tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @alexm-redhat @heheda12345 @ApostaC @orozery
 /tests/weight_loading @mgoin @youkaichao @yewentao256
@@ -171,7 +171,6 @@ mkdocs.yaml @hmellor

 # Pooling models
 /examples/pooling @noooop
-/docs/models/pooling_models @noooop
 /tests/models/*/pooling* @noooop
 /tests/entrypoints/pooling @noooop
 /vllm/config/pooler.py @noooop
@@ -27,7 +27,7 @@ pull_request_rules:
        Hi @{{author}}, the pre-commit checks have failed. Please run:

        ```bash 
-        uv pip install pre-commit>=4.5.1
+        uv pip install pre-commit
        pre-commit install
        pre-commit run --all-files
        ```
@@ -260,7 +260,7 @@ pull_request_rules:
      - files=examples/offline_inference/structured_outputs.py
      - files=examples/online_serving/structured_outputs/structured_outputs.py
      - files~=^tests/v1/structured_output/
-      - files=tests/entrypoints/llm/test_struct_output_generate.py
+      - files=tests/v1/entrypoints/llm/test_struct_output_generate.py
      - files~=^vllm/v1/structured_output/
  actions:
    label:
@@ -333,10 +333,9 @@ pull_request_rules:
    - label != stale
    - or:
      - files~=^tests/tool_use/
-      - files~=^tests/tool_parsers/
-      - files~=^tests/entrypoints/openai/.*tool.*
-      - files~=^tests/entrypoints/anthropic/.*tool.*
-      - files~=^vllm/tool_parsers/
+      - files~=^tests/entrypoints/openai/tool_parsers/
+      - files=tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
+      - files~=^vllm/entrypoints/openai/tool_parsers/
      - files=docs/features/tool_calling.md
      - files~=^examples/tool_chat_*
      - files=examples/offline_inference/chat_with_tools.py
@@ -382,7 +381,7 @@ pull_request_rules:
    - or:
      - files~=^vllm/model_executor/model_loader/tensorizer.py
      - files~=^vllm/model_executor/model_loader/tensorizer_loader.py
-      - files~=^tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+      - files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
      - files~=^tests/model_executor/model_loader/tensorizer_loader/
  actions:
    assign:
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+set -eu
+
+# ensure 1 argument is passed
+if [ "$#" -ne 1 ]; then
+    echo "Usage: $0 <pr_number>"
+    exit 1
+fi
+
+PR_NUMBER=$1
+OLD=/tmp/orig_pr_body.txt
+NEW=/tmp/new_pr_body.txt
+
+gh pr view --json body --template "{{.body}}" "${PR_NUMBER}" > "${OLD}"
+cp "${OLD}" "${NEW}"
+
+# Remove markdown comments (like the <!-- markdownlint-disable --> at the start)
+sed -i '/<!--.*-->$/d' "${NEW}"
+
+# Remove "PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED."
+sed -i '/PLEASE FILL IN THE PR DESCRIPTION HERE.*$/d' "${NEW}"
+
+# Remove all lines after and including "**BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE**"
+sed -i '/\*\*BEFORE SUBMITTING, PLEASE READ.*\*\*/,$d' "${NEW}"
+
+# Remove HTML <details> section that includes <summary> text of "PR Checklist (Click to Expand)"
+python3 - <<EOF
+import regex as re
+
+with open("${NEW}", "r") as file:
+    content = file.read()
+
+pattern = re.compile(r'(---\n\n)?<details>.*?<summary>.*?PR Checklist \(Click to Expand\).*?</summary>.*?</details>', re.DOTALL)
+content = re.sub(pattern, '', content)
+
+with open("${NEW}", "w") as file:
+    file.write(content)
+EOF
+
+# Run this only if ${NEW} is different than ${OLD}
+if ! cmp -s "${OLD}" "${NEW}"; then
+    gh pr edit --body-file "${NEW}" "${PR_NUMBER}"
+    echo
+    echo "Updated PR body:"
+    echo
+    cat "${NEW}"
+else
+    echo "No changes needed"
+fi
@@ -0,0 +1,32 @@
+name: Cleanup PR Body
+
+on:
+  pull_request_target:
+    types: [opened, reopened, edited]
+
+permissions:
+  pull-requests: write
+
+jobs:
+  update-description:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+
+      - name: Set up Python
+        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+
+      - name: Install Python dependencies
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install regex
+
+      - name: Update PR description
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: bash .github/scripts/cleanup_pr_body.sh "${{ github.event.number }}"
@@ -383,107 +383,4 @@ jobs:
                  core.notice(`All users for label "${label}" already mentioned, skipping comment`);
                }
              }
-            }
-
-      - name: Request missing ROCm info from issue author
-        if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
-        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd  # v8.0.0
-        with:
-          script: |
-            const body = (context.payload.issue.body || '').toLowerCase();
-
-            // Check for existing bot comments to avoid duplicate requests
-            const comments = await github.rest.issues.listComments({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              issue_number: context.issue.number,
-            });
-            const botAlreadyAsked = comments.data.some(
-              c => c.user.type === 'Bot' && c.body.includes('<!-- rocm-info-request -->')
-            );
-            if (botAlreadyAsked) {
-              core.notice('ROCm info request already posted, skipping');
-              return;
-            }
-
-            // Define required information and detection patterns
-            const requiredInfo = [
-              {
-                name: 'Reproducer',
-                patterns: [
-                  /reproduc/i, /minimal.?example/i, /repro\b/i, /steps to reproduce/i,
-                  /code.?snippet/i, /sample.?code/i,
-                  /```python[\s\S]*?```/, /```bash[\s\S]*?```/, /```sh[\s\S]*?```/,
-                ],
-                ask: 'A minimal reproducer (code snippet or script that triggers the issue)',
-              },
-              {
-                name: 'Error message',
-                patterns: [
-                  /error/i, /traceback/i, /exception/i, /fault/i, /crash/i,
-                  /failed/i, /abort/i, /panic/i,
-                ],
-                ask: 'The full error message or traceback',
-              },
-              {
-                name: 'Installation method',
-                patterns: [
-                  /docker/i, /rocm\/pytorch/i, /dockerfile/i, /from source/i,
-                  /pip install/i, /build.?from/i, /container/i, /image/i,
-                  /wheel/i, /\.whl/i, /nightly/i,
-                ],
-                ask: 'How you installed vLLM (Docker image name, pip install, or build from source steps)',
-              },
-              {
-                name: 'Command',
-                patterns: [
-                  /vllm serve/i, /python\s+\S+\.py/i, /```bash[\s\S]*?```/,
-                  /```sh[\s\S]*?```/, /command/i, /launch/i, /run\s/i,
-                  /--model/i, /--tensor-parallel/i, /--gpu-memory/i,
-                ],
-                ask: 'The command you used to launch vLLM (e.g., `vllm serve ...` or the Python script)',
-              },
-              {
-                name: 'GFX architecture',
-                patterns: [
-                  /gfx\d{3,4}/i, /mi\d{3}/i, /mi\d{2}\b/i, /radeon/i,
-                  /gpu.?arch/i, /rocm-smi/i, /rocminfo/i, /navi/i,
-                  /instinct/i,
-                ],
-                ask: 'Your GPU model and GFX architecture (e.g., MI300X / gfx942) — run `rocminfo | grep gfx`',
-              },
-            ];
-
-            const issueBody = context.payload.issue.body || '';
-            const missing = requiredInfo.filter(info =>
-              !info.patterns.some(p => p.test(issueBody))
-            );
-
-            if (missing.length === 0) {
-              core.notice('All required ROCm info appears to be present');
-              return;
-            }
-
-            const author = context.payload.issue.user.login;
-            const checklist = requiredInfo.map(info => {
-              const found = !missing.includes(info);
-              return `- [${found ? 'x' : ' '}] ${info.ask}`;
-            }).join('\n');
-            const message = [
-              '<!-- rocm-info-request -->',
-              `Hi @${author}, thanks for reporting this ROCm issue!`,
-              '',
-              'To help us investigate, please make sure the following information is included:',
-              '',
-              checklist,
-              '',
-              'Please provide any unchecked items above. This will help us reproduce and resolve the issue faster. Thank you!',
-            ].join('\n');
-
-            await github.rest.issues.createComment({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              issue_number: context.issue.number,
-              body: message,
-            });
-            core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+            }
@@ -1,9 +1,9 @@
 name: macOS Apple Silicon Smoke Test

 on:
-  schedule:
-    # Daily at 2:30 AM UTC
-    - cron: '30 2 * * *'
+  push:
+    branches:
+      - main
  workflow_dispatch:  # Manual trigger

 permissions:
@@ -1,96 +0,0 @@
-name: New PR Bot
-
-on:
-  pull_request_target:
-    types: [opened]
-
-permissions:
-  pull-requests: write
-
-jobs:
-  update-description:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Update PR description
-        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
-        with:
-          script: |
-            const { owner, repo } = context.repo;
-            const pr_number = context.issue.number;
-
-            const { data: pr } = await github.rest.pulls.get({
-              owner,
-              repo,
-              pull_number: pr_number,
-            });
-
-            let body = pr.body || '';
-            const original = body;
-
-            // Remove markdown comments (<!-- ... -->)
-            body = body.replace(/^<!--.*-->$/gm, '');
-
-            // Remove "PLEASE FILL IN THE PR DESCRIPTION HERE ..."
-            body = body.replace(/^PLEASE FILL IN THE PR DESCRIPTION HERE.*$/gm, '');
-
-            // Remove all lines after and including "**BEFORE SUBMITTING, PLEASE READ ..."
-            body = body.replace(/\*\*BEFORE SUBMITTING, PLEASE READ.*\*\*[\s\S]*$/, '');
-
-            // Remove <details> section containing "PR Checklist (Click to Expand)"
-            body = body.replace(/(---\n\n)?<details>[\s\S]*?<summary>[\s\S]*?PR Checklist \(Click to Expand\)[\s\S]*?<\/summary>[\s\S]*?<\/details>/g, '');
-
-            if (body !== original) {
-              await github.rest.pulls.update({
-                owner,
-                repo,
-                pull_number: pr_number,
-                body,
-              });
-              console.log('Updated PR body');
-            } else {
-              console.log('No changes needed');
-            }
-
-  reminder-comment:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Post welcome comment for first-time contributors
-        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
-        with:
-          script: |
-            const { owner, repo } = context.repo;
-            const prAuthor = context.payload.pull_request.user.login;
-
-            const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
-              q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
-              per_page: 1,
-            });
-
-            const authorPRCount = searchResults.total_count;
-            console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
-
-            if (authorPRCount === 1) {
-              console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
-              await github.rest.issues.createComment({
-                owner,
-                repo,
-                issue_number: context.issue.number,
-                body: [
-                  '\u{1f44b} Hi! Thank you for contributing to the vLLM project.',
-                  '',
-                  '\u{1f4ac} Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.',
-                  '',
-                  'Just a reminder: PRs would not trigger full CI run by default.',
-                  '',
-                  'Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.',
-                  '',
-                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.',
-                  '',
-                  'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.',
-                  '',
-                  '\u{1f680}',
-                ].join('\n'),
-              });
-            } else {
-              console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
-            }
@@ -11,39 +11,9 @@ concurrency:

 permissions:
  contents: read
-  pull-requests: read

 jobs:
-  pre-run-check:
-    if: github.event_name == 'pull_request'
-    runs-on: ubuntu-latest
-    steps:
-    - name: Check PR label and author merge count
-      uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
-      with:
-        script: |
-          const { data: pr } = await github.rest.pulls.get({
-            ...context.repo,
-            pull_number: context.payload.pull_request.number,
-          });
-
-          const hasReadyLabel = pr.labels.some(l => l.name === 'ready');
-
-          const { data: mergedPRs } = await github.rest.search.issuesAndPullRequests({
-            q: `repo:${context.repo.owner}/${context.repo.repo} is:pr is:merged author:${pr.user.login}`,
-            per_page: 4,
-          });
-          const mergedCount = mergedPRs.total_count;
-
-          if (hasReadyLabel || mergedCount >= 4) {
-            core.info(`Check passed: ready label=${hasReadyLabel}, 4+ merged PRs=${mergedCount >= 4}`);
-          } else {
-            core.setFailed(`PR must have the 'ready' label or the author must have at least 4 merged PRs (found ${mergedCount}).`);
-          }
-
  pre-commit:
-    needs: pre-run-check
-    if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
@@ -0,0 +1,54 @@
+name: PR Reminder Comment Bot
+permissions:
+  pull-requests: write
+on:
+  pull_request_target:
+    types: [opened]
+jobs:
+  pr_reminder:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Remind to run full CI on PR
+        uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
+        with:
+          script: |
+            try {
+              // Get the PR author
+              const prAuthor = context.payload.pull_request.user.login;
+              
+              // Check if this is the author's first PR in this repository
+              // Use GitHub's search API to find all PRs by this author
+              const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
+                q: `repo:${context.repo.owner}/${context.repo.repo} type:pr author:${prAuthor}`,
+                per_page: 100  
+              });
+              
+              const authorPRCount = searchResults.total_count;
+              
+              console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
+              
+              // Only post comment if this is the first PR (only one PR by this author)
+              if (authorPRCount === 1) {
+                console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
+                await github.rest.issues.createComment({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                issue_number: context.issue.number,
+                body: '👋 Hi! Thank you for contributing to the vLLM project.\n\n' +
+                  '💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.\n\n' +
+                  'Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run `fastcheck` CI which starts running only a small and essential subset of CI tests to quickly catch errors. \n\n' +
+                  'You ask your reviewers to trigger select CI tests on top of `fastcheck` CI. \n\n' +
+                  'Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.\n\n' +
+                  'To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.\n\n' +
+                  'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.\n\n' +
+                  '🚀'
+                });
+              } else {
+                console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
+              }
+            } catch (error) {
+              console.error('Error checking PR history or posting comment:', error);
+              // Don't fail the workflow, just log the error
+            }
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -108,7 +108,7 @@ uv.lock
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
-.python-version
+# .python-version

 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
@@ -340,9 +340,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")

  list(APPEND VLLM_EXT_SRC
    "csrc/quantization/awq/gemm_kernels.cu"
+    "csrc/permute_cols.cu"
    "csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
    "csrc/quantization/fp4/nvfp4_quant_entry.cu"
    "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
+    "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
    "csrc/cutlass_extensions/common.cpp"
    "csrc/quantization/w8a8/fp8/per_token_group_quant.cu"
    "csrc/quantization/w8a8/int8/per_token_group_quant.cu")
@@ -618,6 +620,31 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

+  #
+  # 2:4 Sparse Kernels
+
+  # The 2:4 sparse kernels cutlass_scaled_sparse_mm and cutlass_compressor
+  # require CUDA 12.2 or later (and only work on Hopper).
+  cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
+    set(SRCS "csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SPARSE_SCALED_MM_C3X=1")
+    message(STATUS "Building sparse_scaled_mm_c3x for archs: ${SCALED_MM_ARCHS}")
+  else()
+    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
+      message(STATUS "Not building sparse_scaled_mm_c3x kernels as CUDA Compiler version is "
+                     "not >= 12.2, we recommend upgrading to CUDA 12.2 or later "
+                     "if you intend on running FP8 sparse quantized models on Hopper.")
+    else()
+      message(STATUS "Not building sparse_scaled_mm_c3x as no compatible archs found "
+                     "in CUDA target architectures")
+    endif()
+  endif()
+
  # The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
  # CUDA 12.8 or later
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
@@ -664,7 +691,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    list(APPEND VLLM_EXT_SRC "${SRCS}")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
-    set(VLLM_NVFP4_SM100_ENABLED TRUE)
    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
  else()
    message(STATUS "Not building NVFP4 as no compatible archs were found.")
@@ -960,54 +986,6 @@ define_extension_target(
 # Setting this variable sidesteps the issue by calling the driver directly.
 target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

-# Propagate ENABLE_NVFP4_SM100 to all languages (including C++ files such as
-# torch_bindings.cpp) so that per-SM op registrations are compiled in.
-if(VLLM_NVFP4_SM100_ENABLED)
-  target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
-endif()
-
-# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
-# https://github.com/vllm-project/vllm/issues/35163 is resolved
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  #
-  # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
-  #
-  set(VLLM_STABLE_EXT_SRC
-    "csrc/libtorch_stable/torch_bindings.cpp")
-
-  if(VLLM_GPU_LANG STREQUAL "CUDA")
-    list(APPEND VLLM_STABLE_EXT_SRC "csrc/libtorch_stable/permute_cols.cu")
-  endif()
-
-  if(VLLM_GPU_LANG STREQUAL "CUDA")
-    set_gencode_flags_for_srcs(
-      SRCS "${VLLM_STABLE_EXT_SRC}"
-      CUDA_ARCHS "${CUDA_ARCHS}")
-  endif()
-
-  message(STATUS "Enabling C_stable extension.")
-  define_extension_target(
-    _C_stable_libtorch
-    DESTINATION vllm
-    LANGUAGE ${VLLM_GPU_LANG}
-    SOURCES ${VLLM_STABLE_EXT_SRC}
-    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-    ARCHITECTURES ${VLLM_GPU_ARCHES}
-    USE_SABI 3
-    WITH_SOABI)
-
-  # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-  # This ensures we only use C-shim APIs available in PyTorch 2.10.
-  # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-  # which is currently set to 2.10.
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    TORCH_TARGET_VERSION=0x020A000000000000ULL)
-
-  # Needed to use cuda APIs from C-shim
-  target_compile_definitions(_C_stable_libtorch PRIVATE
-    USE_CUDA)
-endif()
-
 #
 # _moe_C extension
 #
@@ -1021,7 +999,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/gpt_oss_router_gemm.cu"
    "csrc/moe/router_gemm.cu")
 endif()

@@ -47,8 +47,6 @@ from common import (
    is_mla_backend,
 )

-from vllm.v1.worker.workspace import init_workspace_manager
-

 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
    """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
@@ -464,7 +462,7 @@ def main():
    parser.add_argument(
        "--batch-specs",
        nargs="+",
-        default=None,
+        default=["q2k", "8q1s1k"],
        help="Batch specifications using extended grammar",
    )

@@ -480,21 +478,6 @@ def main():
    parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
    parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
    parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
-    parser.add_argument(
-        "--kv-cache-dtype",
-        default="auto",
-        choices=["auto", "fp8"],
-        help="KV cache dtype: auto or fp8",
-    )
-    parser.add_argument(
-        "--cuda-graphs",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help=(
-            "Launch kernels with CUDA graphs to eliminate CPU overhead"
-            "in measurements (default: True)"
-        ),
-    )

    # Parameter sweep (use YAML config for advanced sweeps)
    parser.add_argument(
@@ -553,24 +536,21 @@ def main():

        # Batch specs and sizes
        # Support both explicit batch_specs and generated batch_spec_ranges
-        # CLI --batch-specs takes precedence over YAML when provided.
-        cli_batch_specs_provided = args.batch_specs is not None
-        if not cli_batch_specs_provided:
-            if "batch_spec_ranges" in yaml_config:
-                # Generate batch specs from ranges
-                generated_specs = generate_batch_specs_from_ranges(
-                    yaml_config["batch_spec_ranges"]
-                )
-                # Combine with any explicit batch_specs
-                if "batch_specs" in yaml_config:
-                    args.batch_specs = yaml_config["batch_specs"] + generated_specs
-                else:
-                    args.batch_specs = generated_specs
-                console.print(
-                    f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
-                )
-            elif "batch_specs" in yaml_config:
-                args.batch_specs = yaml_config["batch_specs"]
+        if "batch_spec_ranges" in yaml_config:
+            # Generate batch specs from ranges
+            generated_specs = generate_batch_specs_from_ranges(
+                yaml_config["batch_spec_ranges"]
+            )
+            # Combine with any explicit batch_specs
+            if "batch_specs" in yaml_config:
+                args.batch_specs = yaml_config["batch_specs"] + generated_specs
+            else:
+                args.batch_specs = generated_specs
+            console.print(
+                f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
+            )
+        elif "batch_specs" in yaml_config:
+            args.batch_specs = yaml_config["batch_specs"]

        if "batch_sizes" in yaml_config:
            args.batch_sizes = yaml_config["batch_sizes"]
@@ -595,10 +575,6 @@ def main():
            args.warmup_iters = yaml_config["warmup_iters"]
        if "profile_memory" in yaml_config:
            args.profile_memory = yaml_config["profile_memory"]
-        if "kv_cache_dtype" in yaml_config:
-            args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
-        if "cuda_graphs" in yaml_config:
-            args.cuda_graphs = yaml_config["cuda_graphs"]

        # Parameter sweep configuration
        if "parameter_sweep" in yaml_config:
@@ -653,18 +629,12 @@ def main():
    # Determine backends
    backends = args.backends or ([args.backend] if args.backend else ["flash"])
    prefill_backends = getattr(args, "prefill_backends", None)
-    if not args.batch_specs:
-        args.batch_specs = ["q2k", "8q1s1k"]
    console.print(f"Backends: {', '.join(backends)}")
    if prefill_backends:
        console.print(f"Prefill backends: {', '.join(prefill_backends)}")
    console.print(f"Batch specs: {', '.join(args.batch_specs)}")
-    console.print(f"KV cache dtype: {args.kv_cache_dtype}")
-    console.print(f"CUDA graphs: {args.cuda_graphs}")
    console.print()

-    init_workspace_manager(args.device)
-
    # Run benchmarks
    all_results = []

@@ -717,8 +687,6 @@ def main():
                        repeats=args.repeats,
                        warmup_iters=args.warmup_iters,
                        profile_memory=args.profile_memory,
-                        kv_cache_dtype=args.kv_cache_dtype,
-                        use_cuda_graphs=args.cuda_graphs,
                    )

                    # Add decode pipeline config
@@ -871,8 +839,6 @@ def main():
            "repeats": args.repeats,
            "warmup_iters": args.warmup_iters,
            "profile_memory": args.profile_memory,
-            "kv_cache_dtype": args.kv_cache_dtype,
-            "use_cuda_graphs": args.cuda_graphs,
        }
        all_results = run_model_parameter_sweep(
            backends,
@@ -895,8 +861,6 @@ def main():
            "repeats": args.repeats,
            "warmup_iters": args.warmup_iters,
            "profile_memory": args.profile_memory,
-            "kv_cache_dtype": args.kv_cache_dtype,
-            "use_cuda_graphs": args.cuda_graphs,
        }
        all_results = run_parameter_sweep(
            backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -927,8 +891,6 @@ def main():
                            repeats=args.repeats,
                            warmup_iters=args.warmup_iters,
                            profile_memory=args.profile_memory,
-                            kv_cache_dtype=args.kv_cache_dtype,
-                            use_cuda_graphs=args.cuda_graphs,
                        )

                        result = run_benchmark(config)
@@ -213,9 +213,6 @@ class BenchmarkConfig:
    profile_memory: bool = False
    use_cuda_graphs: bool = False

-    # "auto" or "fp8"
-    kv_cache_dtype: str = "auto"
-
    # MLA-specific
    prefill_backend: str | None = None
    kv_lora_rank: int | None = None
@@ -372,7 +369,6 @@ class ResultsFormatter:
                    "backend",
                    "batch_spec",
                    "num_layers",
-                    "kv_cache_dtype",
                    "mean_time",
                    "std_time",
                    "throughput",
@@ -386,7 +382,6 @@ class ResultsFormatter:
                        "backend": r.config.backend,
                        "batch_spec": r.config.batch_spec,
                        "num_layers": r.config.num_layers,
-                        "kv_cache_dtype": r.config.kv_cache_dtype,
                        "mean_time": r.mean_time,
                        "std_time": r.std_time,
                        "throughput": r.throughput_tokens_per_sec or 0,
@@ -30,9 +30,9 @@ batch_specs:
  - "2q16k_32q1s4k"         # 2 very large prefill + 32 decode

  # Context extension + decode
-  - "2q1ks2k_16q1s1k"       # 2 extend + 16 decode
-  - "4q2ks4k_32q1s2k"       # 4 extend + 32 decode
-  - "2q1ks8k_32q1s2k"       # 2 large extend + 32 decode
+  - "2q1kkv2k_16q1s1k"       # 2 extend + 16 decode
+  - "4q2kkv4k_32q1s2k"       # 4 extend + 32 decode
+  - "2q1kkv8k_32q1s2k"       # 2 large extend + 32 decode

  # Explicitly chunked prefill
  - "q8k"           # 8k prefill with chunking hint
@@ -1,58 +0,0 @@
-# MLA decode-only benchmark configuration
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128  # Base value, can be swept for TP simulation
-  num_kv_heads: 1  # MLA uses single latent KV
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128  # CUTLASS MLA and FlashAttn MLA use 128
-
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
-batch_specs:
-  # Small batches, varying sequence lengths
-  - "16q1s512"     # 16 requests, 512 KV cache
-  - "16q1s1k"      # 16 requests, 1k KV cache
-  - "16q1s2k"      # 16 requests, 2k KV cache
-  - "16q1s4k"      # 16 requests, 4k KV cache
-
-  # Medium batches
-  - "32q1s1k"      # 32 requests, 1k KV cache
-  - "32q1s2k"      # 32 requests, 2k KV cache
-  - "32q1s4k"      # 32 requests, 4k KV cache
-  - "32q1s8k"      # 32 requests, 8k KV cache
-
-  # Large batches
-  - "64q1s1k"      # 64 requests, 1k KV cache
-  - "64q1s2k"      # 64 requests, 2k KV cache
-  - "64q1s4k"      # 64 requests, 4k KV cache
-  - "64q1s8k"      # 64 requests, 8k KV cache
-
-  # Very large batches
-  - "128q1s1k"     # 128 requests, 1k KV cache
-  - "128q1s2k"     # 128 requests, 2k KV cache
-  - "128q1s4k"     # 128 requests, 4k KV cache
-  - "128q1s8k"     # 128 requests, 8k KV cache
-
-  # Long context
-  - "32q1s16k"     # 32 requests, 16k KV cache
-  - "32q1s32k"     # 32 requests, 32k KV cache
-
-backends:
-  - FLASHMLA_SPARSE
-  - FLASHINFER_MLA_SPARSE
-
-device: "cuda:0"
-repeats: 100
-warmup_iters: 10
-profile_memory: true
@@ -60,11 +60,9 @@ def create_minimal_vllm_config(
    model_name: str = "deepseek-v3",
    block_size: int = 128,
    max_num_seqs: int = 256,
-    max_num_batched_tokens: int = 8192,
    mla_dims: dict | None = None,
    index_topk: int | None = None,
    prefill_backend: str | None = None,
-    kv_cache_dtype: str = "auto",
 ) -> VllmConfig:
    """
    Create minimal VllmConfig for MLA benchmarks.
@@ -151,13 +149,13 @@ def create_minimal_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
-        cache_dtype=kv_cache_dtype,
+        cache_dtype="auto",
        enable_prefix_caching=False,
    )

    scheduler_config = SchedulerConfig(
        max_num_seqs=max_num_seqs,
-        max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
+        max_num_batched_tokens=8192,
        max_model_len=32768,
        is_encoder_decoder=False,
        enable_chunked_prefill=True,
@@ -537,7 +535,6 @@ def _create_backend_impl(
    device: torch.device,
    max_num_tokens: int = 8192,
    index_topk: int | None = None,
-    kv_cache_dtype: str = "auto",
 ):
    """
    Create backend implementation instance.
@@ -586,7 +583,7 @@ def _create_backend_impl(
        "num_kv_heads": mla_dims["num_kv_heads"],
        "alibi_slopes": None,
        "sliding_window": None,
-        "kv_cache_dtype": kv_cache_dtype,
+        "kv_cache_dtype": "auto",
        "logits_soft_cap": None,
        "attn_type": "decoder",
        "kv_sharing_target_layer_name": None,
@@ -704,7 +701,6 @@ def _run_single_benchmark(
    mla_dims: dict,
    device: torch.device,
    indexer=None,
-    kv_cache_dtype: str | None = None,
 ) -> BenchmarkResult:
    """
    Run a single benchmark iteration.
@@ -738,124 +734,49 @@ def _run_single_benchmark(
    )

    # Create KV cache
-    if kv_cache_dtype is None:
-        kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
-    head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
-    if kv_cache_dtype == "fp8_ds_mla":
-        # FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
-        # Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
-        #         + 2*rope_dim bf16 bytes
-        # = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            656,
-            device=device,
-            dtype=torch.uint8,
-        )
-    elif kv_cache_dtype == "fp8":
-        from vllm.platforms import current_platform
+    kv_cache = torch.zeros(
+        num_blocks,
+        block_size,
+        mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
+        device=device,
+        dtype=torch.bfloat16,
+    )

-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            head_size,
-            device=device,
-            dtype=torch.uint8,
-        ).view(current_platform.fp8_dtype())
-    else:
-        kv_cache = torch.zeros(
-            num_blocks,
-            block_size,
-            head_size,
-            device=device,
-            dtype=torch.bfloat16,
-        )
+    # Create input tensors for both decode and prefill modes
+    decode_inputs, prefill_inputs = _create_input_tensors(
+        total_q,
+        mla_dims,
+        backend_cfg["query_format"],
+        device,
+        torch.bfloat16,
+    )

    # Fill indexer with random indices for sparse backends
    is_sparse = backend_cfg.get("is_sparse", False)
    if is_sparse and indexer is not None:
        indexer.fill_random_indices(total_q, max_kv_len)

-    # Determine which forward methods to use based on metadata.
-    # Sparse MLA backends always use forward_mqa
-    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
-    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
-    if not has_decode and not has_prefill:
+    # Determine which forward method to use based on metadata
+    if metadata.decode is not None:
+        forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
+    elif metadata.prefill is not None:
+        forward_fn = lambda: impl.forward_mha(
+            prefill_inputs["q"],
+            prefill_inputs["k_c_normed"],
+            prefill_inputs["k_pe"],
+            kv_cache,
+            metadata,
+            prefill_inputs["k_scale"],
+            prefill_inputs["output"],
+        )
+    else:
        raise RuntimeError("Metadata has neither decode nor prefill metadata")

-    num_decode = (
-        metadata.num_decode_tokens
-        if (has_decode and has_prefill)
-        else total_q
-        if has_decode
-        else 0
-    )
-    num_prefill = total_q - num_decode
-
-    # Some backends requires fp8 queries when using fp8 KV cache.
-    is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
-    quantize_query = is_fp8_kvcache and getattr(
-        impl, "supports_quant_query_input", False
-    )
-
-    # quantize_query forces concat format
-    query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
-
-    # Create decode query tensors
-    if has_decode:
-        decode_inputs, _ = _create_input_tensors(
-            num_decode, mla_dims, query_fmt, device, torch.bfloat16
-        )
-        # Cast decode query to fp8 if the backend supports it
-        if quantize_query:
-            from vllm.platforms import current_platform
-
-            if isinstance(decode_inputs, tuple):
-                decode_inputs = torch.cat(list(decode_inputs), dim=-1)
-            decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
-
-    # Create prefill input tensors
-    if has_prefill:
-        _, prefill_inputs = _create_input_tensors(
-            num_prefill, mla_dims, query_fmt, device, torch.bfloat16
-        )
-
-    # Build forward function
-    def forward_fn():
-        results = []
-        if has_decode:
-            results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
-        if has_prefill:
-            results.append(
-                impl.forward_mha(
-                    prefill_inputs["q"],
-                    prefill_inputs["k_c_normed"],
-                    prefill_inputs["k_pe"],
-                    kv_cache,
-                    metadata,
-                    prefill_inputs["k_scale"],
-                    prefill_inputs["output"],
-                )
-            )
-        return results[0] if len(results) == 1 else tuple(results)
-
    # Warmup
    for _ in range(config.warmup_iters):
        forward_fn()
    torch.accelerator.synchronize()

-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            forward_fn()
-        benchmark_fn = graph.replay
-    else:
-        benchmark_fn = forward_fn
-
    # Benchmark
    times = []
    for _ in range(config.repeats):
@@ -864,7 +785,7 @@ def _run_single_benchmark(

        start.record()
        for _ in range(config.num_layers):
-            benchmark_fn()
+            forward_fn()
        end.record()

        torch.accelerator.synchronize()
@@ -931,30 +852,13 @@ def _run_mla_benchmark_batched(
    # Determine if this is a sparse backend
    is_sparse = backend_cfg.get("is_sparse", False)

-    # Extract kv_cache_dtype from the first config
-    kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
-
-    # FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
-    # Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
-    if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
-        kv_cache_dtype = "fp8_ds_mla"
-
-    # Compute max total_q across all configs so the metadata builder buffer
-    # and scheduler config are large enough for all batch specs.
-    max_total_q = max(
-        sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
-        for cfg, *_ in configs_with_params
-    )
-
    # Create and set vLLM config for MLA (reused across all benchmarks)
    vllm_config = create_minimal_vllm_config(
        model_name="deepseek-v3",  # Used only for model path
        block_size=block_size,
-        max_num_batched_tokens=max_total_q,
        mla_dims=mla_dims,  # Use custom dims from config or default
        index_topk=index_topk if is_sparse else None,
        prefill_backend=prefill_backend,
-        kv_cache_dtype=kv_cache_dtype,
    )

    results = []
@@ -979,9 +883,7 @@ def _run_mla_benchmark_batched(
            mla_dims,
            vllm_config,
            device,
-            max_num_tokens=max_total_q,
            index_topk=index_topk if is_sparse else None,
-            kv_cache_dtype=kv_cache_dtype,
        )

        # Verify the actual prefill backend matches what was requested
@@ -1040,7 +942,6 @@ def _run_mla_benchmark_batched(
                    mla_dims,
                    device,
                    indexer=indexer,
-                    kv_cache_dtype=kv_cache_dtype,
                )
                results.append(result)

@@ -140,7 +140,7 @@ def _create_vllm_config(

    cache_config = CacheConfig(
        block_size=config.block_size,
-        cache_dtype=config.kv_cache_dtype,
+        cache_dtype="auto",
    )
    cache_config.num_gpu_blocks = max_num_blocks
    cache_config.num_cpu_blocks = 0
@@ -215,7 +215,7 @@ def _create_backend_impl(
        num_kv_heads=config.num_kv_heads,
        alibi_slopes=None,
        sliding_window=None,
-        kv_cache_dtype=config.kv_cache_dtype,
+        kv_cache_dtype="auto",
    )

    kv_cache_spec = FullAttentionSpec(
@@ -288,22 +288,12 @@ def _create_input_tensors(
    total_q: int,
    device: torch.device,
    dtype: torch.dtype,
-    quantize_query: bool = False,
 ) -> tuple:
-    """Create Q, K, V input tensors for all layers.
-
-    When quantize_query is True, queries are cast to fp8 to match backends
-    that require query/key/value dtype consistency.
-    """
-    q_dtype = dtype
-    if quantize_query:
-        from vllm.platforms import current_platform
-
-        q_dtype = current_platform.fp8_dtype()
+    """Create Q, K, V input tensors for all layers."""
    q_list = [
        torch.randn(
            total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-        ).to(q_dtype)
+        )
        for _ in range(config.num_layers)
    ]
    k_list = [
@@ -354,17 +344,10 @@ def _create_kv_cache(
    # Compute inverse permutation to get back to logical view
    inv_order = [stride_order.index(i) for i in range(len(stride_order))]

-    # Use fp8 dtype for cache when requested.
-    cache_dtype = dtype
-    if config.kv_cache_dtype == "fp8":
-        from vllm.platforms import current_platform
-
-        cache_dtype = current_platform.fp8_dtype()
-
    cache_list = []
    for _ in range(config.num_layers):
        # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
+        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
        # Permute to logical view
        cache = cache.permute(*inv_order)
        cache_list.append(cache)
@@ -409,37 +392,6 @@ def _run_single_benchmark(
            )
    torch.accelerator.synchronize()

-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-        benchmark_fn = graph.replay
-    else:
-
-        def benchmark_fn():
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-
    # Benchmark
    times = []
    for _ in range(config.repeats):
@@ -447,7 +399,16 @@ def _run_single_benchmark(
        end = torch.cuda.Event(enable_timing=True)

        start.record()
-        benchmark_fn()
+        for i in range(config.num_layers):
+            impl.forward(
+                layer,
+                q_list[i],
+                k_list[i],
+                v_list[i],
+                cache_list[i],
+                attn_metadata,
+                output=out,
+            )
        end.record()

        torch.accelerator.synchronize()
@@ -457,8 +418,8 @@ def _run_single_benchmark(
    mem_stats = {}
    if config.profile_memory:
        mem_stats = {
-            "allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
-            "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
+            "allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
+            "reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
        }

    return times, mem_stats
@@ -541,12 +502,8 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                common_attn_metadata=common_metadata,
            )

-            # Only quantize queries when the impl supports it
-            quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
-                impl, "supports_quant_query_input", False
-            )
            q_list, k_list, v_list = _create_input_tensors(
-                config, total_q, device, dtype, quantize_query=quantize_query
+                config, total_q, device, dtype
            )

            cache_list = _create_kv_cache(
@@ -40,9 +40,9 @@ LLM engine. You can refer to the `vllm.engine.arg_utils.EngineArgs` for more
 details.
 """

+import dataclasses
 import random
 import time
-from dataclasses import fields

 from vllm import LLM, SamplingParams
 from vllm.engine.arg_utils import EngineArgs
@@ -124,7 +124,7 @@ def main(args):

    # Create the LLM engine
    engine_args = EngineArgs.from_cli_args(args)
-    llm = LLM(**{f.name: getattr(engine_args, f.name) for f in fields(engine_args)})
+    llm = LLM(**dataclasses.asdict(engine_args))
    sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len)

    print("------warm up------")
@@ -32,7 +32,6 @@ import dataclasses
 import json
 import random
 import time
-from dataclasses import fields

 from transformers import PreTrainedTokenizerBase

@@ -197,7 +196,7 @@ def main(args):

    engine_args = EngineArgs.from_cli_args(args)

-    llm = LLM(**{f.name: getattr(engine_args, f.name) for f in fields(engine_args)})
+    llm = LLM(**dataclasses.asdict(engine_args))

    sampling_params = SamplingParams(
        temperature=0,
@@ -3,10 +3,10 @@
 """Benchmark offline prioritization."""

 import argparse
+import dataclasses
 import json
 import random
 import time
-from dataclasses import fields

 from transformers import AutoTokenizer, PreTrainedTokenizerBase

@@ -79,7 +79,7 @@ def run_vllm(
 ) -> float:
    from vllm import LLM, SamplingParams

-    llm = LLM(**{f.name: getattr(engine_args, f.name) for f in fields(engine_args)})
+    llm = LLM(**dataclasses.asdict(engine_args))

    assert all(
        llm.llm_engine.model_config.max_model_len >= (request[1] + request[2])
@@ -95,16 +95,13 @@ def create_logits(
 def measure_memory() -> tuple[int, int]:
    """Return (allocated, reserved) memory in bytes."""
    torch.accelerator.synchronize()
-    return (
-        torch.accelerator.memory_allocated(),
-        torch.accelerator.max_memory_allocated(),
-    )
+    return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()


 def reset_memory_stats():
    """Reset peak memory statistics."""
    reset_buffer_cache()
-    torch.accelerator.reset_peak_memory_stats()
+    torch.cuda.reset_peak_memory_stats()
    torch.accelerator.empty_cache()
    gc.collect()

@@ -0,0 +1,517 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import copy
+import itertools
+import pickle as pkl
+import time
+from collections.abc import Callable, Iterable
+
+import torch
+import torch.utils.benchmark as TBenchmark
+from torch.utils.benchmark import Measurement as TMeasurement
+from utils import make_rand_sparse_tensors
+from weight_shapes import WEIGHT_SHAPES
+
+from vllm import _custom_ops as ops
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
+DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
+DEFAULT_TP_SIZES = [1]
+
+
+# bench
+def bench_fn(
+    label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs
+) -> TMeasurement:
+    min_run_time = 1
+
+    globals = {
+        "args": args,
+        "kwargs": kwargs,
+        "fn": fn,
+    }
+    return TBenchmark.Timer(
+        stmt="fn(*args, **kwargs)",
+        globals=globals,
+        label=label,
+        sub_label=sub_label,
+        description=description,
+    ).blocked_autorange(min_run_time=min_run_time)
+
+
+def bench_int8(
+    dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
+) -> Iterable[TMeasurement]:
+    assert dtype == torch.int8
+    b_compressed, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
+    scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+    scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+    bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
+
+    out = ops.cutlass_scaled_sparse_mm(
+        a, b_compressed, e, scale_a, scale_b, torch.bfloat16
+    )
+    out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
+
+    if not torch.allclose(out, out_ref):
+        print("Incorrect results")
+        print(out)
+        print(out_ref)
+    else:
+        print("Correct results")
+
+    timers = []
+    # pytorch impl - bfloat16
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_bf16_bf16_bf16_matmul-no-scales",
+            torch.mm,
+            a.to(dtype=torch.bfloat16),
+            b.to(dtype=torch.bfloat16),
+        )
+    )
+
+    # pytorch impl - float16
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_fp16_fp16_fp16_matmul-no-scales",
+            torch.mm,
+            a.to(dtype=torch.float16),
+            b.to(dtype=torch.float16),
+        )
+    )
+
+    # cutlass impl
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_i8_i8_bf16_scaled_mm",
+            ops.cutlass_scaled_mm,
+            a,
+            b,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+        )
+    )
+
+    # cutlass with bias
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_i8_i8_bf16_scaled_mm_bias",
+            ops.cutlass_scaled_mm,
+            a,
+            b,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+            bias,
+        )
+    )
+
+    # cutlass sparse impl
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_i8_i8_bf16_scaled_sparse_mm",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+        )
+    )
+
+    # cutlass sparse with bias
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_i8_i8_bf16_scaled_sparse_mm_bias",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+            bias,
+        )
+    )
+
+    return timers
+
+
+def bench_fp8(
+    dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
+) -> Iterable[TMeasurement]:
+    assert dtype == torch.float8_e4m3fn
+    b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)
+    scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+    scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+    bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
+
+    out = ops.cutlass_scaled_sparse_mm(
+        a, b_compressed, e, scale_a, scale_b, torch.bfloat16
+    )
+    out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
+
+    if not torch.allclose(out, out_ref):
+        print("Incorrect results")
+        print(out)
+        print(out_ref)
+    else:
+        print("Correct results")
+
+    timers = []
+
+    # pytorch impl w. bf16
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_bf16_bf16_bf16_matmul-no-scales",
+            torch.mm,
+            a.to(dtype=torch.bfloat16, device="cuda"),
+            b.to(dtype=torch.bfloat16, device="cuda"),
+        )
+    )
+
+    # pytorch impl: bf16 output, without fp8 fast accum
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_fp8_fp8_bf16_scaled_mm",
+            torch._scaled_mm,
+            a,
+            b,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            out_dtype=torch.bfloat16,
+        )
+    )
+
+    # pytorch impl: bf16 output, with fp8 fast accum
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",
+            torch._scaled_mm,
+            a,
+            b,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            out_dtype=torch.bfloat16,
+            use_fast_accum=True,
+        )
+    )
+
+    # pytorch impl: fp16 output, without fp8 fast accum
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_fp8_fp8_fp16_scaled_mm",
+            torch._scaled_mm,
+            a,
+            b,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            out_dtype=torch.float16,
+        )
+    )
+
+    # pytorch impl: fp16 output, with fp8 fast accum
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",
+            torch._scaled_mm,
+            a,
+            b,
+            scale_a=scale_a,
+            scale_b=scale_b,
+            out_dtype=torch.float16,
+            use_fast_accum=True,
+        )
+    )
+
+    # cutlass impl: bf16 output
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_fp8_fp8_bf16_scaled_mm",
+            ops.cutlass_scaled_mm,
+            a,
+            b,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+        )
+    )
+
+    # cutlass impl: bf16 output
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_fp8_fp8_bf16_scaled_sparse_mm",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+        )
+    )
+
+    # cutlass impl: fp16 output
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_fp8_fp8_fp16_scaled_sparse_mm",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.float16,
+        )
+    )
+
+    # cutlass impl: bf16 output, with bias
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.bfloat16,
+            bias,
+        )
+    )
+
+    # cutlass impl: fp16 output, with bias
+    timers.append(
+        bench_fn(
+            label,
+            sub_label,
+            "cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",
+            ops.cutlass_scaled_sparse_mm,
+            a,
+            b_compressed,
+            e,
+            scale_a,
+            scale_b,
+            torch.float16,
+            bias.to(dtype=torch.float16),
+        )
+    )
+
+    return timers
+
+
+def bench(
+    dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
+) -> Iterable[TMeasurement]:
+    if dtype == torch.int8:
+        return bench_int8(dtype, m, k, n, label, sub_label)
+    if dtype == torch.float8_e4m3fn:
+        return bench_fp8(dtype, m, k, n, label, sub_label)
+    raise ValueError(
+        f"Unsupported dtype {dtype}: should be one of torch.int8, torch.float8_e4m3fn."
+    )
+
+
+# runner
+def print_timers(timers: Iterable[TMeasurement]):
+    compare = TBenchmark.Compare(timers)
+    compare.print()
+
+
+def run(
+    dtype: torch.dtype, MKNs: Iterable[tuple[int, int, int]]
+) -> Iterable[TMeasurement]:
+    results = []
+    for m, k, n in MKNs:
+        timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm", f"MKN=({m}x{k}x{n})")
+        print_timers(timers)
+        results.extend(timers)
+
+    return results
+
+
+# output makers
+def make_output(
+    data: Iterable[TMeasurement],
+    MKNs: Iterable[tuple[int, int, int]],
+    base_description: str,
+    timestamp=None,
+):
+    print(f"== All Results {base_description} ====")
+    print_timers(data)
+
+    # pickle all the results
+    timestamp = int(time.time()) if timestamp is None else timestamp
+    with open(f"{base_description}-{timestamp}.pkl", "wb") as f:
+        pkl.dump(data, f)
+
+
+# argparse runners
+
+
+def run_square_bench(args):
+    dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment))
+    MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))
+    data = run(args.dtype, MKNs)
+
+    make_output(data, MKNs, f"square_bench-{args.dtype}")
+
+
+def run_range_bench(args):
+    dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))
+    n = len(dim_sizes)
+    Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes
+    Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes
+    Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes
+    MKNs = list(zip(Ms, Ks, Ns))
+    data = run(args.dtype, MKNs)
+
+    make_output(data, MKNs, f"range_bench-{args.dtype}")
+
+
+def run_model_bench(args):
+    print("Benchmarking models:")
+    for i, model in enumerate(args.models):
+        print(f"[{i}]  {model}")
+
+    def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
+        KNs = []
+        for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
+            KN[tp_split_dim] = KN[tp_split_dim] // tp_size
+            KNs.append(KN)
+        return KNs
+
+    model_bench_data = []
+    models_tps = list(itertools.product(args.models, args.tp_sizes))
+    for model, tp_size in models_tps:
+        Ms = args.batch_sizes
+        KNs = model_shapes(model, tp_size)
+        MKNs = []
+        for m in Ms:
+            for k, n in KNs:
+                MKNs.append((m, k, n))
+
+        data = run(args.dtype, MKNs)
+        model_bench_data.append(data)
+
+    # Print all results
+    for data, model_tp in zip(model_bench_data, models_tps):
+        model, tp_size = model_tp
+        print(f"== Results {args.dtype} {model}-TP{tp_size} ====")
+        print_timers(data)
+
+    timestamp = int(time.time())
+
+    all_data = []
+    for d in model_bench_data:
+        all_data.extend(d)
+    # pickle all data
+    with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:
+        pkl.dump(all_data, f)
+
+
+if __name__ == "__main__":
+
+    def to_torch_dtype(dt):
+        if dt == "int8":
+            return torch.int8
+        if dt == "fp8":
+            return torch.float8_e4m3fn
+        raise ValueError("unsupported dtype")
+
+    parser = FlexibleArgumentParser(
+        description="""
+Benchmark Cutlass GEMM.
+
+    To run square GEMMs:
+        python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64
+    
+    To run constant N and K and sweep M:
+        python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384
+    
+    To run dimensions from a model:
+        python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1
+    
+    Output:
+        - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.
+            """,  # noqa: E501
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--dtype",
+        type=to_torch_dtype,
+        required=True,
+        help="Available options are ['int8', 'fp8']",
+    )
+    subparsers = parser.add_subparsers(dest="cmd")
+
+    square_parser = subparsers.add_parser("square_bench")
+    square_parser.add_argument("--dim-start", type=int, required=True)
+    square_parser.add_argument("--dim-end", type=int, required=True)
+    square_parser.add_argument("--dim-increment", type=int, required=True)
+    square_parser.set_defaults(func=run_square_bench)
+
+    range_parser = subparsers.add_parser("range_bench")
+    range_parser.add_argument("--dim-start", type=int, required=True)
+    range_parser.add_argument("--dim-end", type=int, required=True)
+    range_parser.add_argument("--dim-increment", type=int, required=True)
+    range_parser.add_argument("--m-constant", type=int, default=None)
+    range_parser.add_argument("--n-constant", type=int, default=None)
+    range_parser.add_argument("--k-constant", type=int, default=None)
+    range_parser.set_defaults(func=run_range_bench)
+
+    model_parser = subparsers.add_parser("model_bench")
+    model_parser.add_argument(
+        "--models",
+        nargs="+",
+        type=str,
+        default=DEFAULT_MODELS,
+        choices=WEIGHT_SHAPES.keys(),
+    )
+    model_parser.add_argument(
+        "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES
+    )
+    model_parser.add_argument(
+        "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
+    )
+    model_parser.set_defaults(func=run_model_bench)
+
+    args = parser.parse_args()
+    args.func(args)
@@ -5,6 +5,8 @@

 import torch

+import vllm._custom_ops as ops
+

 def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
    finfo = torch.finfo(torch.float8_e4m3fn)
@@ -37,3 +39,49 @@ def make_rand_tensors(
        return to_fp8(a), to_fp8(b)

    raise ValueError("unsupported dtype")
+
+
+def prune_to_2_4(tensor):
+    # Reshape tensor to [N, 4] where N is number of groups of 4
+    original_shape = tensor.shape
+    reshaped = tensor.reshape(-1, 4)
+
+    # Get indices of top 2 absolute values in each group of 4
+    _, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
+
+    # Create binary mask
+    mask = torch.zeros_like(reshaped)
+    mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))
+
+    # Apply mask and reshape back
+    pruned = reshaped * mask
+
+    # Turn all -0.0 to 0.0
+    pruned[pruned == -0.0] = 0.0
+
+    return pruned.reshape(original_shape)
+
+
+def make_rand_sparse_tensors(
+    dtype: torch.dtype, m: int, n: int, k: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    a = torch.randn((m, k), device="cuda") * 5
+    b = torch.randn((n, k), device="cuda").t() * 5
+
+    b = prune_to_2_4(b.t()).t()
+
+    if dtype == torch.int8:
+        a, b = to_int8(a), to_int8(b)
+    elif dtype == torch.float8_e4m3fn:
+        a, b = to_fp8(a), to_fp8(b)
+    elif dtype == torch.float16:
+        a, b = to_fp16(a), to_fp16(b)
+    elif dtype == torch.bfloat16:
+        a, b = to_bf16(a), to_bf16(b)
+    else:
+        raise ValueError("unsupported dtype")
+
+    b_compressed, e = ops.cutlass_sparse_compress(b.t())
+
+    # Compressed B, Metadata, Original A, B
+    return b_compressed, e, a, b
@@ -750,20 +750,17 @@ def get_weight_block_size_safety(config, default_value=None):


 def get_model_params(config):
-    architectures = getattr(config, "architectures", None) or [type(config).__name__]
-    architecture = architectures[0]
-
-    if architecture == "DbrxForCausalLM":
+    if config.architectures[0] == "DbrxForCausalLM":
        E = config.ffn_config.moe_num_experts
        topk = config.ffn_config.moe_top_k
        intermediate_size = config.ffn_config.ffn_hidden_size
        hidden_size = config.hidden_size
-    elif architecture == "JambaForCausalLM":
+    elif config.architectures[0] == "JambaForCausalLM":
        E = config.num_experts
        topk = config.num_experts_per_tok
        intermediate_size = config.intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
+    elif config.architectures[0] in (
        "DeepseekV2ForCausalLM",
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
@@ -777,7 +774,7 @@ def get_model_params(config):
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
+    elif config.architectures[0] in (
        "Qwen2MoeForCausalLM",
        "Qwen3MoeForCausalLM",
        "Qwen3NextForCausalLM",
@@ -786,27 +783,23 @@ def get_model_params(config):
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
        hidden_size = config.hidden_size
-    elif architecture in (
-        "Qwen3VLMoeForConditionalGeneration",
-        "Qwen3_5MoeForConditionalGeneration",
-        "Qwen3_5MoeTextConfig",
-    ):
+    elif config.architectures[0] == "Qwen3VLMoeForConditionalGeneration":
        text_config = config.get_text_config()
        E = text_config.num_experts
        topk = text_config.num_experts_per_tok
        intermediate_size = text_config.moe_intermediate_size
        hidden_size = text_config.hidden_size
-    elif architecture == "HunYuanMoEV1ForCausalLM":
+    elif config.architectures[0] == "HunYuanMoEV1ForCausalLM":
        E = config.num_experts
        topk = config.moe_topk[0]
        intermediate_size = config.moe_intermediate_size[0]
        hidden_size = config.hidden_size
-    elif architecture == "Qwen3OmniMoeForConditionalGeneration":
+    elif config.architectures[0] == "Qwen3OmniMoeForConditionalGeneration":
        E = config.thinker_config.text_config.num_experts
        topk = config.thinker_config.text_config.num_experts_per_tok
        intermediate_size = config.thinker_config.text_config.moe_intermediate_size
        hidden_size = config.thinker_config.text_config.hidden_size
-    elif architecture == "PixtralForConditionalGeneration":
+    elif config.architectures[0] == "PixtralForConditionalGeneration":
        # Pixtral can contain different LLM architectures,
        # recurse to get their parameters
        return get_model_params(config.get_text_config())
@@ -821,23 +814,6 @@ def get_model_params(config):
    return E, topk, intermediate_size, hidden_size


-def resolve_dtype(config) -> torch.dtype:
-    if current_platform.is_rocm():
-        return torch.float16
-
-    dtype = getattr(config, "dtype", None)
-    if dtype is not None:
-        return dtype
-
-    if hasattr(config, "get_text_config"):
-        text_config = config.get_text_config()
-        dtype = getattr(text_config, "dtype", None)
-        if dtype is not None:
-            return dtype
-
-    return torch.bfloat16
-
-
 def get_quantization_group_size(config) -> int | None:
    """Extract the quantization group size from the HF model config.

@@ -885,7 +861,7 @@ def main(args: argparse.Namespace):
    else:
        ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
        shard_intermediate_size = 2 * intermediate_size // args.tp_size
-    dtype = resolve_dtype(config)
+    dtype = torch.float16 if current_platform.is_rocm() else config.dtype
    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
    use_int8_w8a16 = args.dtype == "int8_w8a16"
    use_int4_w4a16 = args.dtype == "int4_w4a16"
@@ -1,669 +0,0 @@
-"""
-Benchmark: SM103 (B300) FP4 Ultra GEMM vs SM100 (B200) NVFP4 GEMM
-===================================================================
-
-This benchmark compares the performance of the SM103-optimized FP4 Ultra
-GEMM kernel against the SM100 NVFP4 GEMM kernel, both running on B300
-hardware.  It also benchmarks the effect of Programmatic Dependent Launch
-(PDL) on the quant->GEMM pipeline, where the GEMM consumer can begin
-before the quant producer finishes.
-
-SM103 kernels use:
-  - K=768 tile (vs K=256 on SM100)
-  - FP4 Ultra MMA (UltraVs16) schedule
-  - NoSmemWarpSpecialized epilogue
-  - Sm103BlockScaledConfig scale factor layout
-
-PDL kernels additionally set:
-  - cudaLaunchAttributeProgrammaticStreamSerialization on quant (producer)
-  - CUTLASS launch_with_pdl=true on GEMM (enables overlap with next kernel)
-
-Usage:
-    python benchmarks/kernels/benchmark_nvfp4_sm103.py [--mode gemm|quant|e2e|pdl|all]
-
-Requirements:
-    - B300 GPU (SM103 / compute capability 10.3)
-    - CUDA >= 12.9
-    - vLLM built with ENABLE_NVFP4_SM100=1 and SM103 support
-"""
-
-import argparse
-from typing import Optional
-
-import torch
-import vllm._C  # noqa: F401 - registers ops into torch.ops._C
-
-# ============================================================================
-# Helpers
-# ============================================================================
-
-
-def round_up(x: int, y: int) -> int:
-    return ((x + y - 1) // y) * y
-
-
-def get_sm_version() -> int:
-    """Return SM version as integer (e.g., 100, 103, 120)."""
-    cap = torch.cuda.get_device_capability()
-    return cap[0] * 10 + cap[1]
-
-
-def create_nvfp4_tensors(
-    m: int, n: int, k: int, dtype: torch.dtype = torch.bfloat16
-) -> dict:
-    """
-    Create synthetic NVFP4 GEMM input tensors (A, B, scales, alpha).
-
-    A: [m, k/2] uint8 (packed FP4)
-    B: [n, k/2] uint8 (packed FP4, column-major)
-    A_sf: [round_up(m,128), round_up(k/16,4)] float8_e4m3fn (SM100 swizzled)
-    B_sf: [round_up(n,128), round_up(k/16,4)] float8_e4m3fn (SM100 swizzled)
-    alpha: [1] float32
-    D: [m, n] output
-    """
-    # Packed FP4 data (random bytes -- content doesn't affect timing)
-    A = torch.randint(0, 256, (m, k // 2), dtype=torch.uint8, device="cuda")
-    B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
-
-    # Scale factors (SM100 swizzled layout)
-    sf_m = round_up(m, 128)
-    sf_n = round_up(n, 128)
-    sf_k = round_up(k // 16, 4)
-
-    A_sf_sm100 = torch.randint(
-        0, 256, (sf_m, sf_k), dtype=torch.uint8, device="cuda"
-    ).view(torch.float8_e4m3fn)
-    B_sf_sm100 = torch.randint(
-        0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
-    ).view(torch.float8_e4m3fn)
-
-    # SM103 layout: convert from SM100 layout
-    A_sf_sm103 = torch.empty_like(A_sf_sm100)
-    B_sf_sm103 = torch.empty_like(B_sf_sm100)
-    torch.ops._C.convert_sf_layout_sm100_to_sm103(A_sf_sm103, A_sf_sm100)
-    torch.ops._C.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
-
-    # Global alpha
-    alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
-
-    # Output
-    D = torch.empty(m, n, dtype=dtype, device="cuda")
-
-    return {
-        "A": A,
-        "B": B,
-        "A_sf_sm100": A_sf_sm100,
-        "B_sf_sm100": B_sf_sm100,
-        "A_sf_sm103": A_sf_sm103,
-        "B_sf_sm103": B_sf_sm103,
-        "alpha": alpha,
-        "D": D,
-    }
-
-
-def create_quant_tensors(
-    m: int, n: int, dtype: torch.dtype = torch.bfloat16
-) -> dict:
-    """Create inputs for activation quantization benchmark."""
-    input_tensor = torch.randn(m, n, dtype=dtype, device="cuda")
-    global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
-    return {"input": input_tensor, "global_scale": global_scale}
-
-
-def bench_fn(
-    fn,
-    warmup: int = 20,
-    iters: int = 100,
-    sync: bool = True,
-) -> float:
-    """Benchmark a function, returning median time in microseconds."""
-    # Warmup
-    for _ in range(warmup):
-        fn()
-    if sync:
-        torch.cuda.synchronize()
-
-    # Timed iterations using CUDA events
-    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
-    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
-
-    for i in range(iters):
-        start_events[i].record()
-        fn()
-        end_events[i].record()
-
-    torch.cuda.synchronize()
-
-    times = [s.elapsed_time(e) * 1000 for s, e in zip(start_events, end_events)]
-    times.sort()
-    # Return median in microseconds
-    return times[len(times) // 2]
-
-
-# ============================================================================
-# GEMM Benchmark
-# ============================================================================
-
-
-def benchmark_gemm(
-    m_sizes: list[int],
-    n: int = 7168,
-    k: int = 7168,
-    dtype: torch.dtype = torch.bfloat16,
-) -> list[dict]:
-    """
-    Benchmark SM100 vs SM103 vs SM103+PDL NVFP4 GEMM kernels side by side.
-
-    PDL on the GEMM sets ProgrammaticStreamSerialization, allowing the NEXT
-    kernel on the stream to overlap with the GEMM's tail.  For isolated GEMM
-    calls (no consumer kernel), the PDL overhead should be near-zero.
-    """
-    vllm_ops = torch.ops._C
-
-    has_sm100a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm100a")
-    has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
-    has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
-
-    if not has_sm100a and not has_sm103a:
-        print("WARNING: Neither sm100a nor sm103a ops are available. "
-              "Rebuild with ENABLE_NVFP4_SM100=1.")
-        return []
-
-    results = []
-
-    for m in m_sizes:
-        tensors = create_nvfp4_tensors(m, n, k, dtype)
-        D = tensors["D"]
-        A, B = tensors["A"], tensors["B"]
-        A_sf_sm100, B_sf_sm100 = tensors["A_sf_sm100"], tensors["B_sf_sm100"]
-        A_sf_sm103, B_sf_sm103 = tensors["A_sf_sm103"], tensors["B_sf_sm103"]
-        alpha = tensors["alpha"]
-
-        flops = 2.0 * m * n * k
-
-        time_sm100: Optional[float] = None
-        time_sm103: Optional[float] = None
-        time_sm103_pdl: Optional[float] = None
-
-        if has_sm100a:
-            def run_sm100():
-                vllm_ops.cutlass_scaled_fp4_mm_sm100a(
-                    D, A, B, A_sf_sm100, B_sf_sm100, alpha
-                )
-            time_sm100 = bench_fn(run_sm100, warmup=20, iters=100)
-
-        if has_sm103a:
-            def run_sm103():
-                vllm_ops.cutlass_scaled_fp4_mm_sm103a(
-                    D, A, B, A_sf_sm103, B_sf_sm103, alpha
-                )
-            time_sm103 = bench_fn(run_sm103, warmup=20, iters=100)
-
-        if has_sm103a_pdl:
-            def run_sm103_pdl():
-                vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
-                    D, A, B, A_sf_sm103, B_sf_sm103, alpha
-                )
-            time_sm103_pdl = bench_fn(run_sm103_pdl, warmup=20, iters=100)
-
-        row: dict = {"M": m, "N": n, "K": k}
-
-        if time_sm100 is not None:
-            row["sm100_us"] = time_sm100
-            row["sm100_tflops"] = flops / (time_sm100 * 1e-6) / 1e12
-
-        if time_sm103 is not None:
-            row["sm103_us"] = time_sm103
-            row["sm103_tflops"] = flops / (time_sm103 * 1e-6) / 1e12
-
-        if time_sm103_pdl is not None:
-            row["sm103pdl_us"] = time_sm103_pdl
-            row["sm103pdl_tflops"] = flops / (time_sm103_pdl * 1e-6) / 1e12
-
-        if time_sm100 is not None and time_sm103 is not None:
-            row["sm103_vs_100"] = time_sm100 / time_sm103
-
-        results.append(row)
-
-    return results
-
-
-# ============================================================================
-# Quantization Benchmark
-# ============================================================================
-
-
-def benchmark_quant(
-    m_sizes: list[int],
-    n: int = 7168,
-    dtype: torch.dtype = torch.bfloat16,
-) -> list[dict]:
-    """
-    Benchmark SM100 vs SM103 activation quantization (BF16 -> NVFP4).
-    """
-    vllm_ops = torch.ops._C
-    results = []
-
-    has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
-
-    for m in m_sizes:
-        tensors = create_quant_tensors(m, n, dtype)
-        input_t = tensors["input"]
-        global_scale = tensors["global_scale"]
-
-        # SM100 quantization (swizzled layout)
-        def run_sm100_quant():
-            vllm_ops.scaled_fp4_quant(input_t, global_scale, True)
-
-        time_sm100 = bench_fn(run_sm100_quant, warmup=20, iters=100)
-
-        row: dict = {
-            "M": m,
-            "N": n,
-            "sm100_us": time_sm100,
-            "sm100_gb_s": (m * n * 2) / (time_sm100 * 1e-6) / 1e9,
-        }
-
-        if has_sm103_quant:
-            def run_sm103_quant():
-                vllm_ops.scaled_fp4_quant_sm103(input_t, global_scale)
-
-            time_sm103 = bench_fn(run_sm103_quant, warmup=20, iters=100)
-            row["sm103_us"] = time_sm103
-            row["sm103_gb_s"] = (m * n * 2) / (time_sm103 * 1e-6) / 1e9
-
-        results.append(row)
-
-    return results
-
-
-# ============================================================================
-# SF Layout Conversion Benchmark
-# ============================================================================
-
-
-def benchmark_sf_conversion(
-    m_sizes: list[int],
-    k: int = 7168,
-) -> list[dict]:
-    """
-    Benchmark the SM100 <-> SM103 scale factor layout conversion kernel.
-
-    This measures the overhead of converting scale factors between layouts,
-    which happens once at model load time for weights.
-    """
-    vllm_ops = torch.ops._C
-    results = []
-
-    for m in m_sizes:
-        sf_m = round_up(m, 128)
-        sf_k = round_up(k // 16, 4)
-
-        # Create source SF tensor (SM100 layout)
-        src = torch.randint(
-            0, 256, (sf_m, sf_k), dtype=torch.uint8, device="cuda"
-        ).view(torch.float8_e4m3fn)
-
-        # Allocate destination (same shape)
-        dst = torch.empty_like(src)
-
-        # Benchmark SM100 -> SM103 conversion
-        def run_convert():
-            vllm_ops.convert_sf_layout_sm100_to_sm103(dst, src)
-
-        time_us = bench_fn(run_convert, warmup=20, iters=200)
-
-        results.append({
-            "M": m,
-            "K": k,
-            "sf_shape": f"{sf_m}x{sf_k}",
-            "kernel": "SM100->SM103 SF convert",
-            "time_us": time_us,
-            "throughput_gb_s": (sf_m * sf_k) / (time_us * 1e-6) / 1e9,
-        })
-
-    return results
-
-
-# ============================================================================
-# End-to-End Benchmark (Quant + GEMM) with PDL comparison
-# ============================================================================
-
-
-def benchmark_e2e(
-    m_sizes: list[int],
-    n: int = 7168,
-    k: int = 7168,
-    dtype: torch.dtype = torch.bfloat16,
-) -> list[dict]:
-    """
-    Benchmark the full NVFP4 inference path: quantize activations + GEMM,
-    comparing SM100, SM103, and SM103+PDL.
-
-    This measures what a real transformer linear layer does:
-    1. Quantize BF16 activations to NVFP4 (with block scales)
-    2. NVFP4 x NVFP4 GEMM
-
-    SM103+PDL enables ProgrammaticStreamSerialization on the quant kernel
-    and launch_with_pdl on the GEMM, allowing the GEMM to begin executing
-    while the quant kernel is still completing its last thread blocks.
-    """
-    vllm_ops = torch.ops._C
-
-    has_sm100a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm100a")
-    has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
-    has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
-    has_sm103_pdl_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103_pdl")
-    has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
-
-    if not has_sm100a and not has_sm103a:
-        print("WARNING: Neither sm100a nor sm103a ops are available. "
-              "Rebuild with ENABLE_NVFP4_SM100=1.")
-        return []
-
-    results = []
-
-    for m in m_sizes:
-        # Create activation input
-        activation = torch.randn(m, k, dtype=dtype, device="cuda")
-        global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
-
-        # Create weight (pre-quantized)
-        B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
-        sf_n = round_up(n, 128)
-        sf_k = round_up(k // 16, 4)
-
-        # Weight SFs in SM100 layout (for SM100 kernel)
-        B_sf_sm100 = torch.randint(
-            0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
-        ).view(torch.float8_e4m3fn)
-        alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
-
-        # Weight SFs in SM103 layout (pre-converted at load time)
-        B_sf_sm103 = torch.empty_like(B_sf_sm100)
-        vllm_ops.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
-
-        D = torch.empty(m, n, dtype=dtype, device="cuda")
-
-        flops = 2.0 * m * n * k
-        row: dict = {"M": m, "N": n, "K": k}
-
-        # --- SM100 baseline: SM100 quant + SM100 GEMM ---
-        if has_sm100a:
-            def run_e2e_sm100():
-                A_q, A_sf = vllm_ops.scaled_fp4_quant(
-                    activation, global_scale, True
-                )
-                A_sf = A_sf.view(torch.float8_e4m3fn)
-                vllm_ops.cutlass_scaled_fp4_mm_sm100a(
-                    D, A_q, B, A_sf, B_sf_sm100, alpha
-                )
-
-            time_sm100 = bench_fn(run_e2e_sm100, warmup=10, iters=50)
-            row["sm100_us"] = time_sm100
-            row["sm100_tflops"] = flops / (time_sm100 * 1e-6) / 1e12
-
-        # --- SM103 without PDL: SM103 quant + SM103 GEMM ---
-        if has_sm103a and has_sm103_quant:
-            def run_e2e_sm103():
-                A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103(
-                    activation, global_scale
-                )
-                A_sf = A_sf.view(torch.float8_e4m3fn)
-                vllm_ops.cutlass_scaled_fp4_mm_sm103a(
-                    D, A_q, B, A_sf, B_sf_sm103, alpha
-                )
-
-            time_sm103 = bench_fn(run_e2e_sm103, warmup=10, iters=50)
-            row["sm103_us"] = time_sm103
-            row["sm103_tflops"] = flops / (time_sm103 * 1e-6) / 1e12
-
-        # --- SM103 with PDL: PDL quant + PDL GEMM ---
-        if has_sm103a_pdl and has_sm103_pdl_quant:
-            def run_e2e_sm103_pdl():
-                # PDL quant: ProgrammaticStreamSerialization allows GEMM to
-                # begin before quant finishes.
-                A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103_pdl(
-                    activation, global_scale
-                )
-                A_sf = A_sf.view(torch.float8_e4m3fn)
-                # PDL GEMM: ProgrammaticStreamSerialization allows the next
-                # layer's kernel to begin before this GEMM finishes.
-                vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
-                    D, A_q, B, A_sf, B_sf_sm103, alpha
-                )
-
-            time_sm103_pdl = bench_fn(run_e2e_sm103_pdl, warmup=10, iters=50)
-            row["sm103pdl_us"] = time_sm103_pdl
-            row["sm103pdl_tflops"] = flops / (time_sm103_pdl * 1e-6) / 1e12
-
-        # Speedup columns
-        if "sm100_us" in row and "sm103_us" in row:
-            row["sm103_vs_100"] = row["sm100_us"] / row["sm103_us"]
-        if "sm103_us" in row and "sm103pdl_us" in row:
-            row["pdl_vs_nop"] = row["sm103_us"] / row["sm103pdl_us"]
-        if "sm100_us" in row and "sm103pdl_us" in row:
-            row["pdl_vs_100"] = row["sm100_us"] / row["sm103pdl_us"]
-
-        results.append(row)
-
-    return results
-
-
-# ============================================================================
-# PDL Pipeline Benchmark (back-to-back quant+GEMM pairs)
-# ============================================================================
-
-
-def benchmark_pdl_pipeline(
-    m_sizes: list[int],
-    n: int = 7168,
-    k: int = 7168,
-    num_layers: int = 4,
-    dtype: torch.dtype = torch.bfloat16,
-) -> list[dict]:
-    """
-    Benchmark the PDL pipeline benefit for back-to-back layers.
-
-    In a real transformer, the same quant->GEMM pattern repeats for each
-    linear layer.  With PDL enabled on both quant and GEMM, each kernel
-    launch overlaps with its predecessor's tail, creating a pipeline:
-
-        quant_1 -> GEMM_1 -> quant_2 -> GEMM_2 -> ...
-
-    This benchmark simulates `num_layers` consecutive quant+GEMM pairs
-    to measure the cumulative pipeline benefit.
-    """
-    vllm_ops = torch.ops._C
-
-    has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
-    has_sm103_pdl_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103_pdl")
-    has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
-    has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
-
-    if not (has_sm103_quant and has_sm103a):
-        print("WARNING: SM103 ops not available.")
-        return []
-
-    results = []
-
-    for m in m_sizes:
-        activation = torch.randn(m, k, dtype=dtype, device="cuda")
-        global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
-        B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
-        sf_n = round_up(n, 128)
-        sf_k = round_up(k // 16, 4)
-        B_sf_sm100 = torch.randint(
-            0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
-        ).view(torch.float8_e4m3fn)
-        B_sf_sm103 = torch.empty_like(B_sf_sm100)
-        vllm_ops.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
-        alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
-        D = torch.empty(m, n, dtype=dtype, device="cuda")
-
-        total_flops = 2.0 * m * n * k * num_layers
-
-        # SM103 without PDL: num_layers sequential quant+GEMM
-        def run_pipeline_no_pdl():
-            for _ in range(num_layers):
-                A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103(
-                    activation, global_scale
-                )
-                A_sf = A_sf.view(torch.float8_e4m3fn)
-                vllm_ops.cutlass_scaled_fp4_mm_sm103a(
-                    D, A_q, B, A_sf, B_sf_sm103, alpha
-                )
-
-        time_no_pdl = bench_fn(run_pipeline_no_pdl, warmup=5, iters=30)
-
-        row: dict = {
-            "M": m, "layers": num_layers,
-            "no_pdl_us": time_no_pdl,
-            "no_pdl_tflops": total_flops / (time_no_pdl * 1e-6) / 1e12,
-        }
-
-        # SM103 with PDL: num_layers pipelined quant+GEMM
-        if has_sm103_pdl_quant and has_sm103a_pdl:
-            def run_pipeline_pdl():
-                for _ in range(num_layers):
-                    A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103_pdl(
-                        activation, global_scale
-                    )
-                    A_sf = A_sf.view(torch.float8_e4m3fn)
-                    vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
-                        D, A_q, B, A_sf, B_sf_sm103, alpha
-                    )
-
-            time_pdl = bench_fn(run_pipeline_pdl, warmup=5, iters=30)
-            row["pdl_us"] = time_pdl
-            row["pdl_tflops"] = total_flops / (time_pdl * 1e-6) / 1e12
-            row["pdl_speedup"] = time_no_pdl / time_pdl
-
-        results.append(row)
-
-    return results
-
-
-# ============================================================================
-# Main
-# ============================================================================
-
-
-def print_results(results: list[dict], title: str):
-    if not results:
-        return
-
-    print(f"\n{'=' * 80}")
-    print(f" {title}")
-    print(f"{'=' * 80}")
-
-    # Determine columns from first result
-    cols = list(results[0].keys())
-    # Header
-    header = " | ".join(f"{c:>15s}" for c in cols)
-    print(header)
-    print("-" * len(header))
-
-    for r in results:
-        row = []
-        for c in cols:
-            v = r.get(c, "")
-            if isinstance(v, float):
-                row.append(f"{v:>15.2f}")
-            elif isinstance(v, int):
-                row.append(f"{v:>15d}")
-            else:
-                row.append(f"{v:>15s}")
-        print(" | ".join(row))
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Benchmark NVFP4 SM103 vs SM100 kernels (with PDL)"
-    )
-    parser.add_argument(
-        "--mode",
-        choices=["gemm", "quant", "sf_convert", "e2e", "pdl", "all"],
-        default="all",
-        help="Which benchmark to run",
-    )
-    parser.add_argument(
-        "--n", type=int, default=7168,
-        help="N dimension (default: 7168, DeepSeek)",
-    )
-    parser.add_argument(
-        "--k", type=int, default=7168,
-        help="K dimension (default: 7168, DeepSeek)",
-    )
-    parser.add_argument(
-        "--layers", type=int, default=4,
-        help="Number of back-to-back layers for PDL pipeline benchmark",
-    )
-    args = parser.parse_args()
-
-    sm = get_sm_version()
-    print(f"GPU: {torch.cuda.get_device_name()}")
-    print(f"SM version: {sm}")
-    print(f"CUDA version: {torch.version.cuda}")
-
-    if sm < 100:
-        print("ERROR: This benchmark requires SM100+ (Blackwell) GPU.")
-        return
-
-    if sm == 103:
-        print("NOTE: Running on SM103 (B300) -- all kernel variants will run.")
-    else:
-        print(f"NOTE: Running on SM{sm} -- SM100 kernel is native; "
-              "SM103 kernel runs via forward compat (may be slower).")
-
-    # Problem sizes typical for LLM inference
-    # Small M = decode, large M = prefill
-    m_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
-
-    if args.mode in ("gemm", "all"):
-        results = benchmark_gemm(m_sizes, n=args.n, k=args.k)
-        print_results(
-            results,
-            f"NVFP4 GEMM: SM100 vs SM103 vs SM103+PDL (N={args.n}, K={args.k})",
-        )
-
-    if args.mode in ("quant", "all"):
-        results = benchmark_quant(m_sizes, n=args.k)
-        print_results(results, f"NVFP4 Activation Quantization (N={args.k})")
-
-    if args.mode in ("sf_convert", "all"):
-        sf_m_sizes = [1024, 2048, 4096, 7168, 8192, 14336, 16384]
-        results = benchmark_sf_conversion(sf_m_sizes, k=args.k)
-        print_results(results, "SF Layout Conversion SM100 <-> SM103")
-
-    if args.mode in ("e2e", "all"):
-        results = benchmark_e2e(m_sizes, n=args.n, k=args.k)
-        print_results(
-            results,
-            f"E2E NVFP4 (Quant+GEMM): SM100 vs SM103 vs SM103+PDL "
-            f"(N={args.n}, K={args.k})",
-        )
-        print(
-            "\nNOTE: sm103_vs_100 = SM100_time / SM103_time (>1 means SM103 faster)\n"
-            "      pdl_vs_nop  = SM103_time / SM103+PDL_time (>1 means PDL faster)\n"
-            "      pdl_vs_100  = SM100_time / SM103+PDL_time (total speedup)"
-        )
-
-    if args.mode in ("pdl", "all"):
-        results = benchmark_pdl_pipeline(
-            m_sizes, n=args.n, k=args.k, num_layers=args.layers
-        )
-        print_results(
-            results,
-            f"PDL Pipeline ({args.layers} layers): SM103 vs SM103+PDL "
-            f"(N={args.n}, K={args.k})",
-        )
-        print(
-            "\nNOTE: pdl_speedup = no_pdl_time / pdl_time\n"
-            "      PDL overlaps quant tail with GEMM head across layer boundaries.\n"
-            "      Benefit is most visible with multiple back-to-back layers."
-        )
-
-
-if __name__ == "__main__":
-    main()
@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.transformers_utils.config import get_config
-from vllm.triton_utils import triton
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# Dimensions supported by the DSV3 specialized kernel
-DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
-DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
-
-# Dimensions supported by the gpt-oss specialized kernel
-GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
-GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
-
-
-def get_batch_size_range(max_batch_size):
-    return [2**x for x in range(14) if 2**x <= max_batch_size]
-
-
-def get_model_params(config):
-    if config.architectures[0] in (
-        "DeepseekV2ForCausalLM",
-        "DeepseekV3ForCausalLM",
-        "DeepseekV32ForCausalLM",
-    ):
-        num_experts = config.n_routed_experts
-        hidden_size = config.hidden_size
-    elif config.architectures[0] in ("GptOssForCausalLM",):
-        num_experts = config.num_local_experts
-        hidden_size = config.hidden_size
-    else:
-        raise ValueError(f"Unsupported architecture: {config.architectures}")
-    return num_experts, hidden_size
-
-
-def get_benchmark(model, max_batch_size, trust_remote_code):
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["batch_size"],
-            x_vals=get_batch_size_range(max_batch_size),
-            x_log=False,
-            line_arg="provider",
-            line_vals=[
-                "torch",
-                "vllm",
-            ],
-            line_names=["PyTorch", "vLLM"],
-            styles=([("blue", "-"), ("red", "-")]),
-            ylabel="TFLOPs",
-            plot_name=f"{model} router gemm throughput",
-            args={},
-        )
-    )
-    def benchmark(batch_size, provider):
-        config = get_config(model=model, trust_remote_code=trust_remote_code)
-        num_experts, hidden_size = get_model_params(config)
-
-        mat_a = torch.randn(
-            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        mat_b = torch.randn(
-            (num_experts, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        bias = torch.randn(
-            num_experts, dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-
-        is_hopper_or_blackwell = current_platform.is_device_capability(
-            90
-        ) or current_platform.is_device_capability_family(100)
-        allow_dsv3_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
-            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_gpt_oss_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
-            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
-        )
-
-        has_bias = False
-        if allow_gpt_oss_router_gemm:
-            has_bias = True
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        if provider == "torch":
-
-            def runner():
-                if has_bias:
-                    F.linear(mat_a, mat_b, bias)
-                else:
-                    F.linear(mat_a, mat_b)
-        elif provider == "vllm":
-
-            def runner():
-                if allow_dsv3_router_gemm:
-                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
-                elif allow_gpt_oss_router_gemm:
-                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
-                else:
-                    raise ValueError("Unsupported router gemm")
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            runner, quantiles=quantiles
-        )
-
-        def tflops(t_ms):
-            flops = 2 * batch_size * hidden_size * num_experts
-            return flops / (t_ms * 1e-3) / 1e12
-
-        return tflops(ms), tflops(max_ms), tflops(min_ms)
-
-    return benchmark
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
-    parser.add_argument("--max-batch-size", default=16, type=int)
-    parser.add_argument("--trust-remote-code", action="store_true")
-    args = parser.parse_args()
-
-    # Get the benchmark function
-    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
-    # Run performance benchmark
-    benchmark.run(print_data=True)
@@ -27,7 +27,7 @@ def get_attn_isa(
    else:
        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
            return "neon"
-        elif torch.cpu._is_amx_tile_supported():
+        elif torch._C._cpu._is_amx_tile_supported():
            return "amx"
        else:
            return "vec"
@@ -24,7 +24,7 @@ except (ImportError, AttributeError) as e:
    sys.exit(1)

 # ISA selection following test_cpu_fused_moe.py pattern
-ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
+ISA_CHOICES = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]


@torch.inference_mode()
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 29210221863736a08f71a866459e368ad1ac4a95
+          GIT_TAG 1488682bb545f7d020e958a33116b1419d1cfc83
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -173,13 +173,10 @@ ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
 void ScratchPadManager::realloc(size_t new_size) {
  new_size = round(new_size);
  if (new_size > size_) {
-    void* new_ptr = std::aligned_alloc(64, new_size);
-    TORCH_CHECK(new_ptr != nullptr,
-                "ScratchPadManager: aligned_alloc failed for size ", new_size);
    if (ptr_ != nullptr) {
      std::free(ptr_);
    }
-    ptr_ = new_ptr;
+    ptr_ = std::aligned_alloc(64, new_size);
    size_ = new_size;
  }
 }
@@ -196,7 +196,6 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
  return val;
 #else
  assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
-  return u32x8_t{};
 #endif
 }

@@ -109,18 +109,16 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,

 #ifndef USE_ROCM
  int flag = 0;
-  CUresult rdma_result = cuDeviceGetAttribute(
+  CUDA_CHECK(cuDeviceGetAttribute(
      &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
-      device);
-  if (rdma_result == CUDA_SUCCESS &&
-      flag) {  // support GPUDirect RDMA if possible
+      device));
+  if (flag) {  // support GPUDirect RDMA if possible
    prop.allocFlags.gpuDirectRDMACapable = 1;
  }
  int fab_flag = 0;
-  CUresult fab_result = cuDeviceGetAttribute(
-      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
-  if (fab_result == CUDA_SUCCESS &&
-      fab_flag) {  // support fabric handle if possible
+  CUDA_CHECK(cuDeviceGetAttribute(
+      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
+  if (fab_flag) {  // support fabric handle if possible
    prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
  }
 #endif
@@ -232,28 +230,6 @@ void unmap_and_release(unsigned long long device, ssize_t size,
    }
  }

-  // ROCm workaround: hipMemRelease does not return physical VRAM to the
-  // free pool while the virtual-address reservation is still held.
-  // Cycling cuMemAddressFree → cuMemAddressReserve (at the same address)
-  // forces the driver to actually release the physical pages while keeping
-  // the same VA available for a later create_and_map.
-  if (first_error == no_error) {
-    first_error = cuMemAddressFree(d_mem, size);
-    if (first_error == no_error) {
-      CUdeviceptr d_mem_new = 0;
-      first_error = cuMemAddressReserve(&d_mem_new, size, 0, d_mem, 0);
-      if (first_error == no_error && d_mem_new != d_mem) {
-        cuMemAddressFree(d_mem_new, size);
-        snprintf(error_msg, sizeof(error_msg),
-                 "ROCm: VA re-reserve got %p instead of %p", (void*)d_mem_new,
-                 (void*)d_mem);
-        error_code = CUresult(1);
-        std::cerr << error_msg << std::endl;
-        return;
-      }
-    }
-  }
-
  if (first_error != no_error) {
    CUDA_CHECK(first_error);
  }
@@ -1,9 +0,0 @@
-#pragma once
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-
-#ifndef USE_ROCM
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& perm);
-#endif
@@ -1,21 +0,0 @@
-#include "ops.h"
-#include "core/registration.h"
-
-#include <torch/csrc/stable/library.h>
-
-// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
-// Note: We register under namespace "_C" so ops are accessible as
-// torch.ops._C.<op_name> for compatibility with existing code.
-STABLE_TORCH_LIBRARY_FRAGMENT(_C, m) {
-#ifndef USE_ROCM
-  m.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-#ifndef USE_ROCM
-  m.impl("permute_cols", TORCH_BOX(&permute_cols));
-#endif
-}
-
-REGISTER_EXTENSION(_C_stable_libtorch)
@@ -1,13 +0,0 @@
-#pragma once
-
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
-#include <cuda_runtime.h>
-
-// Utility to get the current CUDA stream for a given device using stable APIs.
-// Returns a cudaStream_t for use in kernel launches.
-inline cudaStream_t get_current_cuda_stream(int32_t device_index) {
-  void* stream_ptr = nullptr;
-  TORCH_ERROR_CODE_CHECK(
-      aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
-  return reinterpret_cast<cudaStream_t>(stream_ptr);
-}
@@ -1,144 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_cuda.cu
- * Copyright (c) 2025, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
- * All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAStream.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
-#include "gpt_oss_router_gemm.cuh"
-
-void launch_gpt_oss_router_gemm(__nv_bfloat16* gA, __nv_bfloat16* gB,
-                                __nv_bfloat16* gC, __nv_bfloat16* bias,
-                                int batch_size, int output_features,
-                                int input_features, cudaStream_t stream) {
-  static int const WARP_TILE_M = 16;
-  static int const TILE_M = WARP_TILE_M;
-  static int const TILE_N = 8;
-  static int const TILE_K = 64;
-  static int const STAGES = 16;
-  static int const STAGE_UNROLL = 4;
-  static bool const PROFILE = false;
-
-  CUtensorMap weight_map{};
-  CUtensorMap activation_map{};
-
-  constexpr uint32_t rank = 2;
-  uint64_t size[rank] = {(uint64_t)input_features, (uint64_t)output_features};
-  uint64_t stride[rank - 1] = {input_features * sizeof(__nv_bfloat16)};
-  uint32_t box_size[rank] = {TILE_K, TILE_M};
-  uint32_t elem_stride[rank] = {1, 1};
-
-  CUresult res = cuTensorMapEncodeTiled(
-      &weight_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, rank,
-      gB, size, stride, box_size, elem_stride,
-      CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-      CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-      CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-      CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-  TORCH_CHECK(res == CUDA_SUCCESS,
-              "cuTensorMapEncodeTiled failed for weight_map, error code=",
-              static_cast<int>(res));
-
-  size[1] = batch_size;
-  box_size[1] = TILE_N;
-
-  res = cuTensorMapEncodeTiled(
-      &activation_map, CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
-      rank, gA, size, stride, box_size, elem_stride,
-      CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
-      CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B,
-      CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
-      CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-  TORCH_CHECK(res == CUDA_SUCCESS,
-              "cuTensorMapEncodeTiled failed for activation_map, error code=",
-              static_cast<int>(res));
-
-  int smem_size = STAGES * STAGE_UNROLL *
-                  (TILE_M * TILE_K * sizeof(__nv_bfloat16) +
-                   TILE_N * TILE_K * sizeof(__nv_bfloat16));
-
-  gpuErrChk(cudaFuncSetAttribute(
-      gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
-                                 STAGE_UNROLL, PROFILE>,
-      cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-
-  int tiles_m = (output_features + TILE_M - 1) / TILE_M;
-  int tiles_n = (batch_size + TILE_N - 1) / TILE_N;
-
-  dim3 grid(tiles_m, tiles_n);
-  dim3 block(384);
-
-  cudaLaunchConfig_t config;
-  cudaLaunchAttribute attrs[1];
-  config.gridDim = grid;
-  config.blockDim = block;
-  config.dynamicSmemBytes = smem_size;
-  config.stream = stream;
-  config.attrs = attrs;
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.numAttrs = 1;
-
-  cudaLaunchKernelEx(
-      &config,
-      &gpt_oss_router_gemm_kernel<WARP_TILE_M, TILE_M, TILE_N, TILE_K, STAGES,
-                                  STAGE_UNROLL, PROFILE>,
-      gC, gA, gB, bias, output_features, batch_size, input_features, weight_map,
-      activation_map, nullptr);
-}
-
-void gpt_oss_router_gemm_cuda_forward(torch::Tensor& output,
-                                      torch::Tensor input, torch::Tensor weight,
-                                      torch::Tensor bias) {
-  auto const batch_size = input.size(0);
-  auto const input_dim = input.size(1);
-  auto const output_dim = weight.size(0);
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  if (input.scalar_type() == at::ScalarType::BFloat16) {
-    launch_gpt_oss_router_gemm((__nv_bfloat16*)input.data_ptr(),
-                               (__nv_bfloat16*)weight.data_ptr(),
-                               (__nv_bfloat16*)output.mutable_data_ptr(),
-                               (__nv_bfloat16*)bias.data_ptr(), batch_size,
-                               output_dim, input_dim, stream);
-  } else {
-    throw std::invalid_argument("Unsupported dtype, only supports bfloat16");
-  }
-}
-
-void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
-                         torch::Tensor weight, torch::Tensor bias) {
-  TORCH_CHECK(input.dim() == 2, "input must be 2D");
-  TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
-  TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
-  TORCH_CHECK(input.sizes()[1] == weight.sizes()[1],
-              "input.size(1) must match weight.size(1)");
-  TORCH_CHECK(weight.sizes()[0] == bias.sizes()[0],
-              "weight.size(0) must match bias.size(0)");
-  TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16,
-              "input tensor must be bfloat16");
-  TORCH_CHECK(weight.scalar_type() == at::ScalarType::BFloat16,
-              "weight tensor must be bfloat16");
-  TORCH_CHECK(bias.scalar_type() == at::ScalarType::BFloat16,
-              "bias tensor must be bfloat16");
-  gpt_oss_router_gemm_cuda_forward(output, input, weight, bias);
-}
@@ -1,447 +0,0 @@
-/*
- * Adapted from
- * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc7/cpp/tensorrt_llm/kernels/tinygemm2/tinygemm2_kernel.cuh
- * Copyright (c) 2025, The vLLM team.
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
- * All rights reserved. SPDX-License-Identifier: Apache-2.0
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "cuda_bf16.h"
-#include <stdint.h>
-#include <stdio.h>
-#include <vector>
-
-#include "cuda_pipeline.h"
-#include <cuda.h>
-#include <cuda/barrier>
-#include <cuda/std/utility>
-#include <cuda_runtime.h>
-
-using barrier = cuda::barrier<cuda::thread_scope_block>;
-namespace cde = cuda::device::experimental;
-namespace ptx = cuda::ptx;
-
-#define gpuErrChk(ans)                    \
-  {                                       \
-    gpuAssert((ans), __FILE__, __LINE__); \
-  }
-
-inline void gpuAssert(cudaError_t code, char const* file, int line,
-                      bool abort = true) {
-  if (code != cudaSuccess) {
-    fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
-            line);
-    if (abort) {
-      throw std::runtime_error(cudaGetErrorString(code));
-    }
-  }
-}
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-__device__ uint64_t gclock64() {
-  unsigned long long int rv;
-  asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(rv));
-  return rv;
-}
-
-__device__ void ldmatrix(__nv_bfloat16 rv[2], uint32_t smem_ptr) {
-  int dst;
-  asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n"
-               : "=r"(dst)
-               : "r"(smem_ptr));
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = dst;
-}
-
-__device__ void ldmatrix2(__nv_bfloat16 rv[4], uint32_t smem_ptr) {
-  int x, y;
-  asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
-               : "=r"(x), "=r"(y)
-               : "r"(smem_ptr));
-
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = x;
-  rvi[1] = y;
-}
-
-__device__ void ldmatrix4(__nv_bfloat16 rv[8], uint32_t smem_ptr) {
-  int x, y, z, w;
-  asm volatile(
-      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];"
-      : "=r"(x), "=r"(y), "=r"(z), "=r"(w)
-      : "r"(smem_ptr));
-  int* rvi = reinterpret_cast<int*>(&rv[0]);
-  rvi[0] = x;
-  rvi[1] = y;
-  rvi[2] = z;
-  rvi[3] = w;
-}
-
-__device__ void HMMA_1688(float d[4], __nv_bfloat16 a[4], __nv_bfloat16 b[2],
-                          float c[4]) {
-  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
-  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
-  float const* C = reinterpret_cast<float const*>(&c[0]);
-  float* D = reinterpret_cast<float*>(&d[0]);
-
-  asm volatile(
-      "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 "
-      "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
-      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
-      : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]),
-        "f"(C[3]));
-}
-
-__device__ void HMMA_16816(float d[4], __nv_bfloat16 a[8], __nv_bfloat16 b[4],
-                           float c[4]) {
-  uint32_t const* A = reinterpret_cast<uint32_t const*>(&a[0]);
-  uint32_t const* B = reinterpret_cast<uint32_t const*>(&b[0]);
-  float const* C = reinterpret_cast<float const*>(&c[0]);
-  float* D = reinterpret_cast<float*>(&d[0]);
-
-  asm volatile(
-      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
-      : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3])
-      : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]),
-        "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]));
-}
-
-__device__ void bar_wait(uint32_t bar_ptr, int phase) {
-  asm volatile(
-      "{\n"
-      ".reg .pred                P1;\n"
-      "LAB_WAIT:\n"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-      "@P1                       bra.uni DONE;\n"
-      "bra.uni                   LAB_WAIT;\n"
-      "DONE:\n"
-      "}\n" ::"r"(bar_ptr),
-      "r"(phase));
-}
-
-__device__ bool bar_try_wait(uint32_t bar_ptr, int phase) {
-  uint32_t success;
-  #ifdef INTERNAL
-  asm volatile(".pragma \"set knob DontInsertYield\";\n" : : : "memory");
-  #endif
-  asm volatile(
-      "{\n\t"
-      ".reg .pred P1; \n\t"
-      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
-      "selp.b32 %0, 1, 0, P1; \n\t"
-      "}"
-      : "=r"(success)
-      : "r"(bar_ptr), "r"(phase));
-  return success;
-}
-
-__device__ uint32_t elect_one_sync() {
-  uint32_t pred = 0;
-  uint32_t laneid = 0;
-  asm volatile(
-      "{\n"
-      ".reg .b32 %%rx;\n"
-      ".reg .pred %%px;\n"
-      "     elect.sync %%rx|%%px, %2;\n"
-      "@%%px mov.s32 %1, 1;\n"
-      "     mov.s32 %0, %%rx;\n"
-      "}\n"
-      : "+r"(laneid), "+r"(pred)
-      : "r"(0xFFFFFFFF));
-  return pred;
-}
-#endif
-
-struct Profile {
-  uint64_t start;
-  uint64_t weight_load_start;
-  uint64_t act_load_start;
-  uint64_t compute_start;
-  uint64_t complete;
-};
-
-template <int WARP_TILE_M, int TILE_M, int TILE_N, int TILE_K, int STAGES,
-          int STAGE_UNROLL, bool PROFILE>
-__global__ __launch_bounds__(384, 1) void gpt_oss_router_gemm_kernel(
-    __nv_bfloat16* output, __nv_bfloat16* weights, __nv_bfloat16* activations,
-    __nv_bfloat16* bias, int M, int N, int K,
-    const __grid_constant__ CUtensorMap weight_map,
-    const __grid_constant__ CUtensorMap activation_map,
-    Profile* profile = nullptr) {
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-
-  if (PROFILE && threadIdx.x == 0 && blockIdx.y == 0)
-    profile[blockIdx.x].start = gclock64();
-
-  extern __shared__ __align__(128) char smem[];
-
-  __nv_bfloat16* sh_weights = (__nv_bfloat16*)&smem[0];
-  __nv_bfloat16* sh_activations =
-      (__nv_bfloat16*)&smem[STAGES * STAGE_UNROLL * TILE_M * TILE_K *
-                            sizeof(__nv_bfloat16)];
-
-  #pragma nv_diag_suppress static_var_with_dynamic_init
-  __shared__ barrier bar_wt_ready[STAGES];
-  __shared__ barrier bar_act_ready[STAGES];
-  __shared__ barrier bar_data_consumed[STAGES];
-
-  __shared__ float4 reduction_buffer[128];
-
-  __shared__ nv_bfloat16 sh_bias[TILE_M];
-
-  if (threadIdx.x == 0) {
-    for (int i = 0; i < STAGES; i++) {
-      init(&bar_wt_ready[i], 1);
-      init(&bar_act_ready[i], 1);
-      init(&bar_data_consumed[i], 32);
-    }
-    ptx::fence_proxy_async(ptx::space_shared);
-    asm volatile("prefetch.tensormap [%0];"
-                 :
-                 : "l"(reinterpret_cast<uint64_t>(&weight_map))
-                 : "memory");
-    asm volatile("prefetch.tensormap [%0];"
-                 :
-                 : "l"(reinterpret_cast<uint64_t>(&activation_map))
-                 : "memory");
-  }
-  __syncthreads();
-
-  int warp_id = threadIdx.x / 32;
-  int lane_id = threadIdx.x % 32;
-
-  int phase = 0;
-
-  int mib = blockIdx.x * TILE_M;
-  int ni = blockIdx.y * TILE_N;
-
-  float accum[4];
-  for (int i = 0; i < 4; i++) accum[i] = 0.f;
-
-  int const K_LOOPS_DMA =
-      (K + 4 * TILE_K * STAGE_UNROLL - 1) / (4 * (TILE_K * STAGE_UNROLL));
-  int const K_LOOPS_COMPUTE = K_LOOPS_DMA;
-
-  // Data loading thread
-  if (warp_id >= 4 && elect_one_sync()) {
-    int stage = warp_id % 4;
-
-    bool weight_warp = warp_id < 8;
-    if (!weight_warp) {
-      cudaGridDependencySynchronize();
-      cudaTriggerProgrammaticLaunchCompletion();
-    }
-
-    for (int ki = 0; ki < K_LOOPS_DMA; ki++) {
-      int k = (ki * 4 + (warp_id % 4)) * TILE_K * STAGE_UNROLL;
-
-      uint64_t desc_ptr_wt = reinterpret_cast<uint64_t>(&weight_map);
-      uint64_t desc_ptr_act = reinterpret_cast<uint64_t>(&activation_map);
-
-      uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
-      uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
-      int bytes_wt = TILE_M * TILE_K * sizeof(__nv_bfloat16);
-      int bytes_act = TILE_N * TILE_K * sizeof(__nv_bfloat16);
-
-      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
-
-      if (weight_warp)
-        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
-                     :
-                     : "r"(bar_ptr_wt), "r"(STAGE_UNROLL * bytes_wt));
-      if (!weight_warp)
-        asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"
-                     :
-                     : "r"(bar_ptr_act), "r"(STAGE_UNROLL * bytes_act));
-
-      if (PROFILE && blockIdx.y == 0 && ki == 0 && weight_warp)
-        profile[blockIdx.x].weight_load_start = gclock64();
-      if (PROFILE && blockIdx.y == 0 && ki == 0 && !weight_warp)
-        profile[blockIdx.x].act_load_start = gclock64();
-
-      for (int i = 0; i < STAGE_UNROLL; i++) {
-        uint32_t smem_ptr_wt = __cvta_generic_to_shared(
-            &sh_weights[(stage * STAGE_UNROLL + i) * TILE_M * TILE_K]);
-        uint32_t crd0 = k + i * TILE_K;
-        uint32_t crd1 = mib;
-        if (weight_warp)
-          asm volatile(
-              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
-              "tx::bytes [%0], [%1, {%3,%4}], "
-              "[%2];"
-              :
-              : "r"(smem_ptr_wt), "l"(desc_ptr_wt), "r"(bar_ptr_wt), "r"(crd0),
-                "r"(crd1)
-              : "memory");
-
-        uint32_t smem_ptr_act = __cvta_generic_to_shared(
-            &sh_activations[(stage * STAGE_UNROLL + i) * TILE_N * TILE_K]);
-        crd0 = k + i * TILE_K;
-        crd1 = ni;
-        if (!weight_warp)
-          asm volatile(
-              "cp.async.bulk.tensor.2d.shared::cta.global.mbarrier::complete_"
-              "tx::bytes [%0], [%1, {%3,%4}], "
-              "[%2];"
-              :
-              : "r"(smem_ptr_act), "l"(desc_ptr_act), "r"(bar_ptr_act),
-                "r"(crd0), "r"(crd1)
-              : "memory");
-      }
-
-      stage += 4;
-      if (stage >= STAGES) {
-        stage = warp_id % 4;
-        phase ^= 1;
-      }
-    }
-    // Wait for pending loads to be consumed before exiting, to avoid race
-    for (int i = 0; i < (STAGES / 4) - 1; i++) {
-      bar_wait(__cvta_generic_to_shared(&bar_data_consumed[stage]), phase ^ 1);
-      stage += 4;
-      if (stage >= STAGES) {
-        stage = warp_id % 4;
-        phase ^= 1;
-      }
-    }
-  }
-  // Compute threads
-  else if (warp_id < 4) {
-    // Sneak the bias load into the compute warps since they're just waiting for
-    // stuff anyway
-    if (threadIdx.x < TILE_M) sh_bias[threadIdx.x] = bias[mib + threadIdx.x];
-
-    int stage = warp_id;
-
-    int phase = 0;
-    int lane_id_div8 = lane_id / 8;
-    int lane_id_mod8 = lane_id % 8;
-
-    int lane_row_offset_wt = (lane_id_div8 % 2) ? 8 : 0;
-    int lane_col_offset_wt = (lane_id_div8 / 2) ? 1 : 0;
-
-    int row_wt = lane_id_mod8 + lane_row_offset_wt;
-    int row_act = lane_id_mod8;
-
-    int row_offset_wt = (reinterpret_cast<uintptr_t>(sh_weights) / 128) % 8;
-    int row_offset_act = row_offset_wt;
-
-    uint32_t bar_ptr_wt = __cvta_generic_to_shared(&bar_wt_ready[stage]);
-    uint32_t bar_ptr_act = __cvta_generic_to_shared(&bar_act_ready[stage]);
-
-    bool weight_ready = bar_try_wait(bar_ptr_wt, phase);
-    bool act_ready = bar_try_wait(bar_ptr_act, phase);
-
-  #pragma unroll 2
-    for (int ki = 0; ki < K_LOOPS_COMPUTE; ki++) {
-      int next_stage = stage + 4;
-      int next_phase = phase;
-      if (next_stage >= STAGES) {
-        next_stage = warp_id;
-        next_phase ^= 1;
-      }
-
-      while (!weight_ready || !act_ready) {
-        weight_ready = bar_try_wait(bar_ptr_wt, phase);
-        act_ready = bar_try_wait(bar_ptr_act, phase);
-      }
-
-      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0 && ki == 0)
-        profile[blockIdx.x].compute_start = gclock64();
-
-      if (ki + 1 < K_LOOPS_COMPUTE) {
-        weight_ready = bar_try_wait(
-            __cvta_generic_to_shared(&bar_wt_ready[next_stage]), next_phase);
-        act_ready = bar_try_wait(
-            __cvta_generic_to_shared(&bar_act_ready[next_stage]), next_phase);
-      }
-
-  #pragma unroll
-      for (int su = 0; su < STAGE_UNROLL; su++) {
-        __nv_bfloat16* ptr_weights =
-            &sh_weights[(stage * STAGE_UNROLL + su) * TILE_M * TILE_K];
-        __nv_bfloat16* ptr_act =
-            &sh_activations[(stage * STAGE_UNROLL + su) * TILE_N * TILE_K];
-
-  #pragma unroll
-        for (int kii = 0; kii < TILE_K / 16; kii++) {
-          __nv_bfloat16 a[8];
-          __nv_bfloat16 b[4];
-
-          int col = 2 * kii + lane_col_offset_wt;
-          int col_sw = ((row_wt + row_offset_wt) % 8) ^ col;
-
-          ldmatrix4(a, __cvta_generic_to_shared(
-                           &ptr_weights[row_wt * TILE_K + col_sw * 8]));
-
-          col = 2 * kii + lane_id_div8;
-          col_sw = ((row_act + row_offset_act) % 8) ^ col;
-
-          ldmatrix2(b, __cvta_generic_to_shared(
-                           &ptr_act[row_act * TILE_K + 8 * col_sw]));
-
-          HMMA_16816(accum, a, b, accum);
-        }
-      }
-
-      uint32_t bar_c = __cvta_generic_to_shared(&bar_data_consumed[stage]);
-      asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" : : "r"(bar_c));
-
-      stage = next_stage;
-      phase = next_phase;
-    }
-
-    float4 accum4;
-    accum4.x = accum[0];
-    accum4.y = accum[1];
-    accum4.z = accum[2];
-    accum4.w = accum[3];
-    reduction_buffer[threadIdx.x] = accum4;
-
-    __syncthreads();
-
-    if (warp_id == 0) {
-      int mi = mib + warp_id * WARP_TILE_M;
-      int tm = mi + lane_id / 4;
-      int tn = ni + 2 * (lane_id % 4);
-
-      float4 accum1 = reduction_buffer[32 + threadIdx.x];
-      float4 accum2 = reduction_buffer[64 + threadIdx.x];
-      float4 accum3 = reduction_buffer[96 + threadIdx.x];
-
-      accum[0] = accum[0] + accum1.x + accum2.x + accum3.x;
-      accum[1] = accum[1] + accum1.y + accum2.y + accum3.y;
-      accum[2] = accum[2] + accum1.z + accum2.z + accum3.z;
-      accum[3] = accum[3] + accum1.w + accum2.w + accum3.w;
-
-      float bias_lo = __bfloat162float(sh_bias[tm - mib]);
-      float bias_hi = __bfloat162float(sh_bias[tm + 8 - mib]);
-
-      if (tn < N && tm < M)
-        output[tn * M + tm] = __float2bfloat16(accum[0] + bias_lo);
-      if (tn + 1 < N && tm < M)
-        output[(tn + 1) * M + tm] = __float2bfloat16(accum[1] + bias_lo);
-      if (tn < N && tm + 8 < M)
-        output[tn * M + tm + 8] = __float2bfloat16(accum[2] + bias_hi);
-      if (tn + 1 < N && tm + 8 < M)
-        output[(tn + 1) * M + tm + 8] = __float2bfloat16(accum[3] + bias_hi);
-
-      if (PROFILE && blockIdx.y == 0 && threadIdx.x == 0)
-        profile[blockIdx.x].complete = gclock64();
-    }
-  }
-#endif  // end if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-}
@@ -70,8 +70,4 @@ torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
                      const torch::Tensor& mat_b);
-
-// gpt-oss optimized router GEMM kernel for SM90+
-void gpt_oss_router_gemm(torch::Tensor& output, torch::Tensor input,
-                         torch::Tensor weight, torch::Tensor bias);
 #endif
@@ -73,9 +73,10 @@ void moe_permute(
  MOE_DISPATCH(input.scalar_type(), [&] {
    expandInputRowsKernelLauncher<scalar_t>(
        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
-        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
-        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
-        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
+        get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
+        get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
+        get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
+        n_hidden, topk, n_local_expert, stream);
  });
 }

@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,

 template <typename T>
 void expandInputRowsKernelLauncher(
-    T const* unpermuted_input, T* permuted_output,
+    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
    int const* expanded_dest_row_to_expanded_source_row,
    int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
    int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -2,7 +2,7 @@

 template <typename T, bool CHECK_SKIPPED>
 __global__ void expandInputRowsKernel(
-    T const* unpermuted_input, T* permuted_output,
+    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
    int const* expanded_dest_row_to_expanded_source_row,
    int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
    int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -16,6 +16,7 @@ __global__ void expandInputRowsKernel(
  int64_t expanded_dest_row = blockIdx.x;
  int64_t const expanded_source_row =
      expanded_dest_row_to_expanded_source_row[expanded_dest_row];
+  int expert_id = sorted_experts[expanded_dest_row];

  if (threadIdx.x == 0) {
    assert(expanded_dest_row <= INT32_MAX);
@@ -53,7 +54,7 @@ __global__ void expandInputRowsKernel(

 template <typename T>
 void expandInputRowsKernelLauncher(
-    T const* unpermuted_input, T* permuted_output,
+    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
    int const* expanded_dest_row_to_expanded_source_row,
    int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
    int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -69,12 +70,12 @@ void expandInputRowsKernelLauncher(
  bool is_check_skip = num_valid_tokens_ptr != nullptr;
  auto func = func_map[is_check_skip];

-  func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
-                                       expanded_dest_row_to_expanded_source_row,
-                                       expanded_source_row_to_expanded_dest_row,
-                                       permuted_idx, expert_first_token_offset,
-                                       num_rows, num_valid_tokens_ptr, cols, k,
-                                       num_local_experts);
+  func<<<blocks, threads, 0, stream>>>(
+      unpermuted_input, permuted_output, sorted_experts,
+      expanded_dest_row_to_expanded_source_row,
+      expanded_source_row_to_expanded_dest_row, permuted_idx,
+      expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
+      num_local_experts);
 }

 template <class T, class U>
@@ -132,12 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  // DeepSeek V3 optimized router GEMM for SM90+
  m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
  // conditionally compiled so impl registration is in source file
-
-  // gpt-oss optimized router GEMM kernel for SM90+
-  m.def(
-      "gpt_oss_router_gemm(Tensor! output, Tensor input, Tensor weights, "
-      "Tensor bias) -> ()");
-  m.impl("gpt_oss_router_gemm", torch::kCUDA, &gpt_oss_router_gemm);
 #endif
 }

@@ -201,6 +201,7 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
                             torch::Tensor _zeros, int64_t split_k_iters,
                             int64_t thx, int64_t thy);

+torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm);
 #endif

 torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
@@ -237,7 +238,6 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
                           torch::Tensor const& B_sf,
                           torch::Tensor const& alpha);

-
 void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
                       torch::Tensor const& b, torch::Tensor const& a_scales,
                       torch::Tensor const& b_scales,
@@ -262,8 +262,7 @@ void get_cutlass_moe_mm_data(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const std::optional<torch::Tensor>& blockscale_offsets);

 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
    const torch::Tensor& expert_first_token_offset,
@@ -286,14 +285,20 @@ void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
                           std::optional<torch::Tensor> const& azp,
                           std::optional<torch::Tensor> const& bias);

-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
-    torch::Tensor const& input, torch::Tensor const& input_scale,
-    bool is_sf_swizzled_layout);
+bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability);

-void scaled_fp4_quant_out(torch::Tensor const& input,
-                          torch::Tensor const& input_scale,
-                          bool is_sf_swizzled_layout, torch::Tensor& output,
-                          torch::Tensor& output_scale);
+void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
+                              torch::Tensor const& b, torch::Tensor const& e,
+                              torch::Tensor const& a_scales,
+                              torch::Tensor const& b_scales,
+                              std::optional<torch::Tensor> const& bias);
+
+std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
+
+void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
+                      torch::Tensor& output_scale,
+                      torch::Tensor const& input_scale,
+                      bool is_sf_swizzled_layout);

 void scaled_fp4_experts_quant(
    torch::Tensor& output, torch::Tensor& output_scale,
@@ -307,11 +312,6 @@ void silu_and_mul_scaled_fp4_experts_quant(
    torch::Tensor const& input_offset_by_experts,
    torch::Tensor const& output_scale_offset_by_experts);

-void convert_sf_layout_sm100_to_sm103(torch::Tensor& dst,
-                                      torch::Tensor const& src);
-void convert_sf_layout_sm103_to_sm100(torch::Tensor& dst,
-                                      torch::Tensor const& src);
-
 void per_token_group_quant_fp8(const torch::Tensor& input,
                               torch::Tensor& output_q, torch::Tensor& output_s,
                               int64_t group_size, double eps, double fp8_min,
@@ -1,13 +1,10 @@
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
+#include <torch/all.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>

 #include <cuda_fp16.h>

-#include "torch_utils.h"
-
 static constexpr int default_threads = 256;
 static constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }

@@ -67,22 +64,19 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,

 // More efficient version of A[..., perm]
 //  taken from gptq_marlin.cu
-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& perm) {
-  const int32_t dev = A.get_device_index();
-  const torch::stable::accelerator::DeviceGuard device_guard(dev);
-  const auto stream = get_current_cuda_stream(dev);
+torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm) {
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
+  auto dev = A.get_device();
+  auto stream = at::cuda::getCurrentCUDAStream(dev);

-  STD_TORCH_CHECK(
-      A.scalar_type() == torch::headeronly::ScalarType::Half ||
-          A.scalar_type() == torch::headeronly::ScalarType::BFloat16,
-      "Currently only 16bit types are supported");
-  STD_TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
-  STD_TORCH_CHECK(A.size(-1) % 8 == 0,
-                  "A columns must be a multiple of 8 (128bits)");
-  auto A_2d = torch::stable::view(A, {-1, A.size(-1)});
+  TORCH_CHECK(A.scalar_type() == at::kHalf || A.scalar_type() == at::kBFloat16,
+              "Currently only 16bit types are supported");
+  TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
+  TORCH_CHECK(A.size(-1) % 8 == 0,
+              "A columns must be a multiple of 8 (128bits)");
+  auto A_2d = A.view({-1, A.size(-1)});

-  torch::stable::Tensor D = torch::stable::empty_like(A);
+  torch::Tensor D = torch::empty_like(A);
  int sms;
  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
  int block_rows = div_ceil(A_2d.size(0), sms);
@@ -16,8 +16,6 @@

 #include <torch/all.h>

-#include "nvfp4_utils.cuh"
-
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
@@ -27,18 +25,6 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
                             bool is_sf_swizzled_layout);
 #endif

-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-void scaled_fp4_quant_sm103a(torch::Tensor const& output,
-                             torch::Tensor const& input,
-                             torch::Tensor const& output_sf,
-                             torch::Tensor const& input_sf);
-// PDL variant: launches quant with ProgrammaticStreamSerialization.
-void scaled_fp4_quant_sm103a_pdl(torch::Tensor const& output,
-                                 torch::Tensor const& input,
-                                 torch::Tensor const& output_sf,
-                                 torch::Tensor const& input_sf);
-#endif
-
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
 void scaled_fp4_experts_quant_sm1xxa(
@@ -65,10 +51,9 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
    torch::Tensor const& output_scale_offset_by_experts);
 #endif

-void scaled_fp4_quant_out(torch::Tensor const& input,
-                          torch::Tensor const& input_sf,
-                          bool is_sf_swizzled_layout, torch::Tensor& output,
-                          torch::Tensor& output_sf) {
+void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
+                      torch::Tensor& output_sf, torch::Tensor const& input_sf,
+                      bool is_sf_swizzled_layout) {
 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
  return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
@@ -77,34 +62,6 @@ void scaled_fp4_quant_out(torch::Tensor const& input,
  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
 }

-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf,
-    bool is_sf_swizzled_layout) {
-  int64_t n = input.size(-1);
-  int64_t m = input.numel() / n;
-  auto device = input.device();
-
-  // Two fp4 values packed into a uint8
-  auto output = torch::empty(
-      {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
-
-  torch::Tensor output_sf;
-  if (is_sf_swizzled_layout) {
-    auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
-    output_sf = torch::empty(
-        {sf_m, sf_n},
-        torch::TensorOptions().device(device).dtype(torch::kInt32));
-  } else {
-    output_sf = torch::empty(
-        {m, n / CVT_FP4_SF_VEC_SIZE},
-        torch::TensorOptions().device(device).dtype(torch::kUInt8));
-  }
-
-  scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
-                       output_sf);
-  return {output, output_sf};
-}
-
 void scaled_fp4_experts_quant(
    torch::Tensor& output, torch::Tensor& output_scale,
    torch::Tensor const& input, torch::Tensor const& input_global_scale,
@@ -144,79 +101,3 @@ void silu_and_mul_scaled_fp4_experts_quant(
  TORCH_CHECK_NOT_IMPLEMENTED(
      false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
 }
-
-// SM103-native quantization: writes SM103-layout scale factors directly,
-// eliminating the SM100->SM103 conversion step on the critical path.
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf) {
-  int64_t n = input.size(-1);
-  int64_t m = input.numel() / n;
-  auto device = input.device();
-
-  auto output = torch::empty(
-      {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
-
-  auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
-  auto output_sf = torch::empty(
-      {sf_m, sf_n},
-      torch::TensorOptions().device(device).dtype(torch::kInt32));
-
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  scaled_fp4_quant_sm103a(output, input, output_sf, input_sf);
-  return {output, output_sf};
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false,
-                              "No compiled SM103 nvfp4 quantization kernel");
-}
-
-void scaled_fp4_quant_sm103a_out(torch::Tensor const& input,
-                                 torch::Tensor const& input_sf,
-                                 torch::Tensor& output,
-                                 torch::Tensor& output_sf) {
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  scaled_fp4_quant_sm103a(output, input, output_sf, input_sf);
-  return;
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(false,
-                              "No compiled SM103 nvfp4 quantization kernel");
-}
-
-// ============================================================================
-// PDL-enabled SM103 quantization entry points.
-//
-// These launch the quant kernel with ProgrammaticStreamSerialization,
-// allowing the subsequent GEMM to begin before quantization completes.
-// ============================================================================
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_pdl_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf) {
-  int64_t n = input.size(-1);
-  int64_t m = input.numel() / n;
-  auto device = input.device();
-
-  auto output = torch::empty(
-      {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
-
-  auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
-  auto output_sf = torch::empty(
-      {sf_m, sf_n},
-      torch::TensorOptions().device(device).dtype(torch::kInt32));
-
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  scaled_fp4_quant_sm103a_pdl(output, input, output_sf, input_sf);
-  return {output, output_sf};
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      false, "No compiled SM103 PDL nvfp4 quantization kernel");
-}
-
-void scaled_fp4_quant_sm103a_pdl_out(torch::Tensor const& input,
-                                     torch::Tensor const& input_sf,
-                                     torch::Tensor& output,
-                                     torch::Tensor& output_sf) {
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  scaled_fp4_quant_sm103a_pdl(output, input, output_sf, input_sf);
-  return;
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      false, "No compiled SM103 PDL nvfp4 quantization kernel");
-}
@@ -171,305 +171,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
  }
 }

-// ============================================================================
-// SM103 (B300) activation quantization kernel.
-//
-// Identical to the SM100 cvt_fp16_to_fp4 except it writes scale factors
-// in the SM103 swizzled layout (Sm103BlockScaledConfig).
-// ============================================================================
-template <class Type, bool UE8M0_SF = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
-    cvt_fp16_to_fp4_sm103(int32_t numRows, int32_t numCols,
-                          int32_t num_padded_cols,
-                          Type const* __restrict__ in,
-                          float const* __restrict__ SFScale,
-                          uint32_t* __restrict__ out,
-                          uint32_t* __restrict__ SFout) {
-  using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;
-
-  static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
-      (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
-  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
-                "Vec size is not matched.");
-
-  int32_t const numKTiles = (numCols + 63) / 64;
-
-  int sf_m = round_up<int>(numRows, 128);
-  int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
-  int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
-
-  float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
-
-  for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
-    if (colIdx < num_padded_cols) {
-      PackedVec in_vec;
-      int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
-
-      bool valid = (rowIdx < numRows) && (elem_idx < numCols);
-      if constexpr (CVT_FP4_PACK16) {
-        ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
-                         &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
-                         valid);
-      } else {
-        ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
-                         &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
-                         valid);
-      }
-
-      // SM103: Use SM103-specific SF offset function
-      auto sf_out =
-          cvt_quant_to_fp4_get_sf_out_offset_sm103<uint32_t,
-                                                   CVT_FP4_NUM_THREADS_PER_SF>(
-              rowIdx, colIdx, numKTiles, SFout);
-
-      auto out_val =
-          cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
-              in_vec, global_scale, sf_out);
-
-      if (valid) {
-        if constexpr (CVT_FP4_PACK16) {
-          int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
-          uint64_t packed64 =
-              (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
-          reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
-        } else {
-          out[inOffset] = out_val;
-        }
-      }
-    }
-  }
-}
-
-// ============================================================================
-// Scale factor layout conversion: SM100 <-> SM103
-//
-// Converts an already-swizzled SF tensor between SM100 and SM103 layouts.
-// Both layouts use the same 512-byte tile structure (128 M-rows x 4 K-cols)
-// but arrange bytes differently within each tile.
-//
-// SM100 offset: outerM(=mIdx%32)*16 + innerM(=(mIdx/32)%4)*4 + innerK
-// SM103 offset: m8(=(mIdx/16)%8)*16 + m4a(=(mIdx/4)%4)*128 + m4b(=mIdx%4)*4
-//               + innerK
-// ============================================================================
-__global__ void convert_sf_sm100_to_sm103_kernel(
-    const uint8_t* __restrict__ src,
-    uint8_t* __restrict__ dst,
-    int32_t numMTiles,
-    int32_t numKTiles) {
-  // Each thread converts one byte (one SF value).
-  // Grid: numMTiles * numKTiles blocks, 512 threads per block.
-  int32_t tile_idx = blockIdx.x;
-  int32_t mTileIdx = tile_idx / numKTiles;
-  int32_t kTileIdx = tile_idx % numKTiles;
-
-  // Each tile is 512 bytes: 128 M-positions x 4 K-positions.
-  int32_t local_idx = threadIdx.x;  // 0..511
-  if (mTileIdx >= numMTiles) return;
-
-  int64_t tile_base = static_cast<int64_t>(tile_idx) << 9;
-
-  // Decode this thread's (mLocal, kLocal) from a simple linear index.
-  int32_t mLocal = local_idx >> 2;   // 0..127
-  int32_t kLocal = local_idx & 3;    // 0..3
-
-  // Compute SM100 source offset within tile.
-  int32_t outerMIdx = mLocal & 31;
-  int32_t innerMIdx = (mLocal >> 5) & 3;
-  int32_t sm100_off = (outerMIdx << 4) | (innerMIdx << 2) | kLocal;
-
-  // Compute SM103 destination offset within tile.
-  int32_t m4b = mLocal & 3;
-  int32_t m4a = (mLocal >> 2) & 3;
-  int32_t m8  = (mLocal >> 4) & 7;
-  int32_t sm103_off = (m8 << 4) | (m4a << 7) | (m4b << 2) | kLocal;
-
-  dst[tile_base + sm103_off] = src[tile_base + sm100_off];
-}
-
-__global__ void convert_sf_sm103_to_sm100_kernel(
-    const uint8_t* __restrict__ src,
-    uint8_t* __restrict__ dst,
-    int32_t numMTiles,
-    int32_t numKTiles) {
-  int32_t tile_idx = blockIdx.x;
-  int32_t mTileIdx = tile_idx / numKTiles;
-  if (mTileIdx >= numMTiles) return;
-
-  int32_t local_idx = threadIdx.x;
-  int64_t tile_base = static_cast<int64_t>(tile_idx) << 9;
-
-  int32_t mLocal = local_idx >> 2;
-  int32_t kLocal = local_idx & 3;
-
-  // SM103 source offset
-  int32_t m4b = mLocal & 3;
-  int32_t m4a = (mLocal >> 2) & 3;
-  int32_t m8  = (mLocal >> 4) & 7;
-  int32_t sm103_off = (m8 << 4) | (m4a << 7) | (m4b << 2) | kLocal;
-
-  // SM100 destination offset
-  int32_t outerMIdx = mLocal & 31;
-  int32_t innerMIdx = (mLocal >> 5) & 3;
-  int32_t sm100_off = (outerMIdx << 4) | (innerMIdx << 2) | kLocal;
-
-  dst[tile_base + sm100_off] = src[tile_base + sm103_off];
-}
-
 }  // namespace vllm

-// ============================================================================
-// Host entry: SM103 activation quantization
-//
-// When use_pdl=true, the kernel is launched with
-// cudaLaunchAttributeProgrammaticStreamSerialization, allowing the next
-// kernel on the same stream (typically the GEMM consumer) to begin
-// executing before this quantization kernel fully completes.  This
-// overlaps the tail of quantization with the head of the GEMM.
-// ============================================================================
-static void scaled_fp4_quant_sm103a_impl(torch::Tensor const& output,
-                                         torch::Tensor const& input,
-                                         torch::Tensor const& output_sf,
-                                         torch::Tensor const& input_sf,
-                                         bool use_pdl) {
-  int32_t m = input.size(0);
-  int32_t n = input.size(1);
-
-  TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
-  TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
-                  input.scalar_type() == at::ScalarType::BFloat16,
-              "Unsupported input data type for quantize_to_fp4.");
-
-  int multiProcessorCount =
-      get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
-
-  auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
-  auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
-  auto output_ptr = static_cast<int64_t*>(output.data_ptr());
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
-
-  int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
-
-  dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
-
-  // SM103 always uses swizzled layout (the SM103 variant)
-  int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
-  int32_t num_padded_cols =
-      sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
-
-  int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
-  int grid_x =
-      std::min(vllm::computeEffectiveRows(m),
-               std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
-  dim3 grid(grid_x, grid_y);
-
-  VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_sm103", [&] {
-    using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
-    auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
-    auto output_u32 = reinterpret_cast<uint32_t*>(output_ptr);
-    auto sf_out_u32 = reinterpret_cast<uint32_t*>(sf_out);
-
-    if (use_pdl) {
-      // PDL launch: set ProgrammaticStreamSerialization so the next kernel
-      // (GEMM) can begin before this quant kernel fully completes.
-      cudaLaunchConfig_t launch_config = {};
-      launch_config.gridDim = grid;
-      launch_config.blockDim = block;
-      launch_config.dynamicSmemBytes = 0;
-      launch_config.stream = stream;
-
-      cudaLaunchAttribute pdl_attr;
-      pdl_attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
-      pdl_attr.val.programmaticStreamSerializationAllowed = 1;
-      launch_config.numAttrs = 1;
-      launch_config.attrs = &pdl_attr;
-
-      CUDA_CHECK(cudaLaunchKernelEx(
-          &launch_config,
-          vllm::cvt_fp16_to_fp4_sm103<cuda_type, false>,
-          m, n, num_padded_cols, input_ptr, input_sf_ptr,
-          output_u32, sf_out_u32));
-    } else {
-      vllm::cvt_fp16_to_fp4_sm103<cuda_type, false>
-          <<<grid, block, 0, stream>>>(
-              m, n, num_padded_cols, input_ptr, input_sf_ptr,
-              output_u32, sf_out_u32);
-    }
-  });
-}
-
-// Original entry point (no PDL).
-void scaled_fp4_quant_sm103a(torch::Tensor const& output,
-                             torch::Tensor const& input,
-                             torch::Tensor const& output_sf,
-                             torch::Tensor const& input_sf) {
-  scaled_fp4_quant_sm103a_impl(output, input, output_sf, input_sf,
-                                /*use_pdl=*/false);
-}
-
-// PDL-enabled entry point: launches quant kernel with
-// ProgrammaticStreamSerialization to overlap with a subsequent GEMM.
-void scaled_fp4_quant_sm103a_pdl(torch::Tensor const& output,
-                                 torch::Tensor const& input,
-                                 torch::Tensor const& output_sf,
-                                 torch::Tensor const& input_sf) {
-  scaled_fp4_quant_sm103a_impl(output, input, output_sf, input_sf,
-                                /*use_pdl=*/true);
-}
-
-// ============================================================================
-// Host entry: SF layout conversion SM100 <-> SM103
-// ============================================================================
-void convert_sf_layout_sm100_to_sm103(torch::Tensor& dst,
-                                      torch::Tensor const& src) {
-  TORCH_CHECK(src.is_contiguous(), "Source SF tensor must be contiguous");
-  TORCH_CHECK(dst.is_contiguous(), "Destination SF tensor must be contiguous");
-  TORCH_CHECK(src.numel() == dst.numel(),
-              "Source and destination must have the same number of elements");
-
-  // SF tensors are stored as int32 with shape (rounded_m, rounded_k / 4)
-  // Total bytes = rounded_m * (rounded_k / 4) * 4 = rounded_m * rounded_k
-  int64_t total_bytes = src.numel() * src.element_size();
-  int32_t numMTiles = src.size(0) / 128;
-  int32_t numKTiles = total_bytes / (numMTiles * 512);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(src));
-  auto stream = at::cuda::getCurrentCUDAStream(src.get_device());
-
-  int32_t num_tiles = numMTiles * numKTiles;
-  dim3 grid(num_tiles);
-  dim3 block(512);
-
-  vllm::convert_sf_sm100_to_sm103_kernel<<<grid, block, 0, stream>>>(
-      static_cast<const uint8_t*>(src.data_ptr()),
-      static_cast<uint8_t*>(dst.data_ptr()),
-      numMTiles, numKTiles);
-}
-
-void convert_sf_layout_sm103_to_sm100(torch::Tensor& dst,
-                                      torch::Tensor const& src) {
-  TORCH_CHECK(src.is_contiguous() && dst.is_contiguous());
-  TORCH_CHECK(src.numel() == dst.numel());
-
-  int64_t total_bytes = src.numel() * src.element_size();
-  int32_t numMTiles = src.size(0) / 128;
-  int32_t numKTiles = total_bytes / (numMTiles * 512);
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(src));
-  auto stream = at::cuda::getCurrentCUDAStream(src.get_device());
-
-  int32_t num_tiles = numMTiles * numKTiles;
-  vllm::convert_sf_sm103_to_sm100_kernel<<<dim3(num_tiles), dim3(512), 0, stream>>>(
-      static_cast<const uint8_t*>(src.data_ptr()),
-      static_cast<uint8_t*>(dst.data_ptr()),
-      numMTiles, numKTiles);
-}
-
-// ============================================================================
-// Original SM100 host entry
-// ============================================================================
 void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
                             torch::Tensor const& input,
                             torch::Tensor const& output_sf,
@@ -24,20 +24,6 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
                                  torch::Tensor const& A_sf,
                                  torch::Tensor const& B_sf,
                                  torch::Tensor const& alpha);
-// SM103 (B300) uses FP4 Ultra MMA -- separate entry point compiled from
-// the same source file, guarded by CUTLASS_ARCH_MMA_SM103_SUPPORTED.
-void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha);
-// PDL variant: GEMM launched with ProgrammaticStreamSerialization.
-void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D,
-                                      torch::Tensor const& A,
-                                      torch::Tensor const& B,
-                                      torch::Tensor const& A_sf,
-                                      torch::Tensor const& B_sf,
-                                      torch::Tensor const& alpha);
 #endif

 #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
@@ -57,14 +43,6 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
  const int32_t sm = get_sm_version_num();

 #if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  // SM103 (B300): Use FP4 Ultra kernels with K=768 tiles for higher
-  // throughput. Falls through to SM100 path if SM103 kernels weren’t compiled
-  // (e.g., CUDA < 12.9).
-  if (sm == 103) {
-    cutlass_scaled_fp4_mm_sm103a(D, A, B, A_sf, B_sf, alpha);
-    return;
-  }
-
  if (sm >= 100 && sm < 120) {
    cutlass_scaled_fp4_mm_sm100a(D, A, B, A_sf, B_sf, alpha);
    return;
@@ -36,10 +36,6 @@ using namespace cute;

 #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)

-// ============================================================================
-// SM100 (B200) Tile Configurations
-// ============================================================================
-
 // Configuration for M in (256, inf)
 struct sm100_fp4_config_default {
  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
@@ -67,51 +63,6 @@ struct sm100_fp4_config_M16 {
  using PerSmTileShape_MNK = Shape<_128, _128, _256>;
 };

-// ============================================================================
-// SM103 (B300 / Blackwell Ultra) Tile Configurations
-//
-// Key differences from SM100:
-//   - Tile K = 768 is MANDATORY (CUTLASS static_assert)
-//   - Uses FP4 Ultra MMA instructions (UltraVs16) for higher throughput
-//   - Uses NoSmem epilogue (saves shared memory for mainloop)
-//   - 1SM for small M, 2SM for large M (cooperative SM pairs)
-// ============================================================================
-#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
-
-// SM103 configuration for M in (256, inf) -- 2SM cooperative execution
-struct sm103_fp4_config_default {
-  // 2SM schedule: two SMs cooperate on one tile for higher throughput
-  using KernelSchedule = cutlass::gemm::
-      KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103;
-  using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized2Sm;
-  using TileShape = Shape<_256, _256, Int<768>>;
-  using ClusterShape = Shape<_2, _2, _1>;
-  using PerSmTileShape_MNK = Shape<_128, _256, Int<768>>;
-};
-
-// SM103 configuration for M in (16, 256] -- 2SM with smaller N tile
-struct sm103_fp4_config_M256 {
-  using KernelSchedule = cutlass::gemm::
-      KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103;
-  using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized2Sm;
-  using TileShape = Shape<_256, _128, Int<768>>;
-  using ClusterShape = Shape<_2, _1, _1>;
-  using PerSmTileShape_MNK = Shape<_128, _128, Int<768>>;
-};
-
-// SM103 configuration for M in [1, 16] -- 1SM (decode / small batch)
-struct sm103_fp4_config_M16 {
-  // 1SM schedule: single SM per tile, lower latency for small problems
-  using KernelSchedule = cutlass::gemm::
-      KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103;
-  using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized1Sm;
-  using TileShape = Shape<_128, _128, Int<768>>;
-  using ClusterShape = Shape<_1, _1, _1>;
-  using PerSmTileShape_MNK = Shape<_128, _128, Int<768>>;
-};
-
-#endif  // CUTLASS_ARCH_MMA_SM103_SUPPORTED
-
 template <typename Config, typename OutType>
 struct Fp4GemmSm100 {
  // A matrix configuration
@@ -174,99 +125,6 @@ struct Fp4GemmSm100 {
  using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
 };

-// ============================================================================
-// SM103 GEMM Definition (FP4 Ultra)
-//
-// SM103 differs from SM100 in several fundamental ways:
-//   1. Uses cutlass::arch::Sm103 (separate CollectiveBuilder specialization)
-//   2. Element types passed as cute::tuple<DataType, ScaleFactorType>
-//      (SM100 uses nv_float4_t<float_e2m1_t> wrapper instead)
-//   3. Tile K = 768 (SM100 uses K = 256)
-//   4. Epilogue uses NoSmemWarpSpecialized (SM100 uses TmaWarpSpecialized)
-//   5. Scale factor memory layout uses Sm103BlockScaledConfig
-//      (different swizzle pattern from SM100's Sm1xxBlockScaledConfig)
-//
-// IMPORTANT: Scale factor layout compatibility
-//   SM103 and SM100 use DIFFERENT physical scale factor layouts in memory.
-//   The activation quantization kernel (scaled_fp4_quant) and the weight
-//   scale factors in NVFP4 checkpoints must produce/store data in the
-//   SM103-expected layout when using these kernels. Passing SM100-format
-//   scale factors to SM103 kernels will produce incorrect results.
-//   See Sm103BlockScaledConfig::tile_atom_to_shape_SFA for the expected
-//   layout.
-// ============================================================================
-#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
-
-template <typename Config, typename OutType>
-struct Fp4GemmSm103 {
-  // A matrix configuration -- bare float_e2m1_t (not nv_float4_t wrapper)
-  using ElementA = cutlass::float_e2m1_t;
-  using ElementSFA = cutlass::float_ue4m3_t;
-  using LayoutATag = cutlass::layout::RowMajor;
-  static constexpr int AlignmentA = 32;
-
-  // B matrix configuration
-  using ElementB = cutlass::float_e2m1_t;
-  using ElementSFB = cutlass::float_ue4m3_t;
-  using LayoutBTag = cutlass::layout::ColumnMajor;
-  static constexpr int AlignmentB = 32;
-
-  // C/D matrix configuration
-  using ElementD = OutType;
-  using ElementC = OutType;
-  using LayoutCTag = cutlass::layout::RowMajor;
-  using LayoutDTag = cutlass::layout::RowMajor;
-  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-
-  // Kernel functional config
-  using ElementAccumulator = float;
-  using ArchTag = cutlass::arch::Sm103;
-  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-
-  // Use config's tile shapes (K=768 mandatory for SM103)
-  using MmaTileShape = typename Config::TileShape;
-  using ClusterShape = typename Config::ClusterShape;
-  using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK;
-
-  // Epilogue: SM103 uses NoSmem variant with OpClassTensorOp
-  // Note: epilogue builder uses Sm100 arch tag (shared epilogue HW)
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
-          PerSmTileShape_MNK, ClusterShape,
-          cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
-          ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
-          LayoutDTag, AlignmentD,
-          typename Config::EpilogueSchedule>::CollectiveOp;
-
-  // Mainloop: SM103 passes element+SF types as tuples to the builder
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, cute::tuple<ElementA, ElementSFA>, LayoutATag,
-          AlignmentA, cute::tuple<ElementB, ElementSFB>, LayoutBTag, AlignmentB,
-          ElementAccumulator, MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename Config::KernelSchedule>::CollectiveOp;
-
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  using StrideA = typename Gemm::GemmKernel::StrideA;
-  using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
-  using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
-  using StrideB = typename Gemm::GemmKernel::StrideB;
-  using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{}));
-  using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
-  using StrideC = typename Gemm::GemmKernel::StrideC;
-  using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{}));
-  using StrideD = typename Gemm::GemmKernel::StrideD;
-  using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
-};
-
-#endif  // CUTLASS_ARCH_MMA_SM103_SUPPORTED
-
 template <typename Config>
 typename Config::Gemm::Arguments args_from_options(
    at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
@@ -319,7 +177,7 @@ template <typename Config>
 void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
             at::Tensor const& A_sf, at::Tensor const& B_sf,
             at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
-             cudaStream_t stream, bool launch_with_pdl = false) {
+             cudaStream_t stream) {
  typename Config::Gemm gemm;

  auto arguments =
@@ -334,12 +192,7 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,

  CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));

-  // When launch_with_pdl=true, CUTLASS sets
-  // cudaLaunchAttributeProgrammaticStreamSerialization on the GEMM kernel,
-  // allowing the next kernel on the stream to begin before this GEMM
-  // fully completes.
-  CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream,
-                         /*cuda_adapter=*/nullptr, launch_with_pdl));
+  CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
 }

 // Dispatch function to select appropriate config based on M
@@ -367,39 +220,6 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
  }
 }

-// ============================================================================
-// SM103 Dispatch
-// ============================================================================
-#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
-
-template <typename OutType>
-void cutlass_fp4_gemm_sm103_dispatch(torch::Tensor& D, torch::Tensor const& A,
-                                     torch::Tensor const& B,
-                                     torch::Tensor const& A_sf,
-                                     torch::Tensor const& B_sf,
-                                     torch::Tensor const& alpha, int64_t m,
-                                     int64_t n, int64_t k,
-                                     cudaStream_t stream,
-                                     bool launch_with_pdl = false) {
-  uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
-
-  if (mp2 <= 16) {
-    // m in [1, 16] -- 1SM, low-latency decode
-    runGemm<Fp4GemmSm103<sm103_fp4_config_M16, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
-  } else if (mp2 <= 256) {
-    // m in (16, 256] -- 2SM, small tile
-    runGemm<Fp4GemmSm103<sm103_fp4_config_M256, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
-  } else {
-    // m in (256, inf) -- 2SM, large tile
-    runGemm<Fp4GemmSm103<sm103_fp4_config_default, OutType>>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
-  }
-}
-
-#endif  // CUTLASS_ARCH_MMA_SM103_SUPPORTED
-
 #else
 template <typename OutType>
 void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
@@ -495,107 +315,3 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
                ")");
  }
 }
-
-// ============================================================================
-// SM103 Entry Point (B300 / Blackwell Ultra)
-//
-// Uses FP4 Ultra MMA instructions with K=768 tiles for higher throughput.
-// Scale factors must be in Sm103BlockScaledConfig layout (different from SM100).
-//
-// When launch_with_pdl=true, the CUTLASS GEMM is launched with
-// ProgrammaticStreamSerialization, allowing the next kernel on the stream
-// to begin before this GEMM completes.  Combined with a PDL-enabled
-// quantization producer, this creates a pipelined quant->GEMM overlap.
-// ============================================================================
-#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
-
-static void cutlass_scaled_fp4_mm_sm103a_impl(
-    torch::Tensor& D, torch::Tensor const& A, torch::Tensor const& B,
-    torch::Tensor const& A_sf, torch::Tensor const& B_sf,
-    torch::Tensor const& alpha, bool launch_with_pdl) {
-  CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
-  CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
-
-  CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
-  CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
-
-  CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
-
-  TORCH_CHECK(A.dim() == 2, "a must be a matrix");
-  TORCH_CHECK(B.dim() == 2, "b must be a matrix");
-  TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
-              "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
-              A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
-
-  auto const m = A.sizes()[0];
-  auto const n = B.sizes()[0];
-  auto const k = A.sizes()[1] * 2;
-
-  constexpr int alignment = 32;
-  TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
-              ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
-              "), k: ", k, ".");
-  TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
-              ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
-
-  // SM103 scale factor shape validation.
-  // Physical dimensions are the same as SM100 (padded to 128 x ceil(k/16,4)),
-  // but the internal swizzle pattern (Sm103BlockScaledConfig) differs.
-  auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
-  int rounded_m = round_up(m, 128);
-  int rounded_n = round_up(n, 128);
-  int rounded_k = round_up(k / 16, 4);
-
-  TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
-  TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
-  TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
-              "scale_a and scale_b shapes cannot be multiplied (",
-              A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
-              "x", B_sf.sizes()[1], ")");
-  TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
-              "scale_a must be padded and swizzled to a shape (", rounded_m,
-              "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
-              A_sf.sizes()[1], ")");
-  TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
-              "scale_b must be padded and swizzled to a shape (", rounded_n,
-              "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
-              B_sf.sizes()[1], ")");
-
-  auto out_dtype = D.dtype();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
-
-  if (out_dtype == at::ScalarType::Half) {
-    cutlass_fp4_gemm_sm103_dispatch<cutlass::half_t>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
-  } else if (out_dtype == at::ScalarType::BFloat16) {
-    cutlass_fp4_gemm_sm103_dispatch<cutlass::bfloat16_t>(
-        D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
-  } else {
-    TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype,
-                ")");
-  }
-}
-
-// Original entry point (no PDL).
-void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha) {
-  cutlass_scaled_fp4_mm_sm103a_impl(D, A, B, A_sf, B_sf, alpha,
-                                     /*launch_with_pdl=*/false);
-}
-
-// PDL-enabled entry point: GEMM launched with ProgrammaticStreamSerialization
-// so the next kernel on the stream can overlap with this GEMM's tail.
-void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D, torch::Tensor const& A,
-                                      torch::Tensor const& B,
-                                      torch::Tensor const& A_sf,
-                                      torch::Tensor const& B_sf,
-                                      torch::Tensor const& alpha) {
-  cutlass_scaled_fp4_mm_sm103a_impl(D, A, B, A_sf, B_sf, alpha,
-                                     /*launch_with_pdl=*/true);
-}
-
-#endif  // CUTLASS_ARCH_MMA_SM103_SUPPORTED
@@ -18,7 +18,6 @@

 #include <cuda_runtime.h>
 #include <cuda_fp8.h>
-#include <utility>

 #include "../../cuda_vec_utils.cuh"

@@ -55,18 +54,6 @@ inline int computeEffectiveRows(int m) {
  return round_up(m, ROW_TILE);
 }

-// Compute the shape of the swizzled SF output tensor.
-// Returns (rounded_m, rounded_n / 4) where:
-//   rounded_m = round_up(m, 128)
-//   rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
-inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
-                                                          int64_t n) {
-  int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
-  int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
-  int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
-  return {rounded_m, rounded_n / 4};
-}
-
 // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
 inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
  uint32_t val;
@@ -199,55 +186,6 @@ __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
  return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
 }

-// ============================================================================
-// SM103 (Blackwell Ultra / B300) swizzled SF offset.
-//
-// SM103 uses Sm103BlockScaledConfig with a 3-level M decomposition:
-//   M -> (m8, m4a, m4b)  where mIdx = m4b + m4a*4 + m8*16
-//   K -> (sfv16_broadcast, k4)
-//
-// Atom layout:
-//   Shape:  <Shape<_8, _4, _4>, Shape<SFVecSize=16, _4>>
-//   Stride: <Stride<_16, _128, _4>, Stride<_0, _1>>
-//
-// Physical offset = m8*16 + m4a*128 + m4b*4 + k4
-// Each 128-row x 4-col tile occupies 512 bytes (same as SM100).
-// ============================================================================
-template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
-__device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset_sm103(
-    int rowIdx, int colIdx, int32_t numKTiles, SFType* SFout) {
-  static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
-                CVT_FP4_NUM_THREADS_PER_SF == 2);
-
-  if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF != 0) {
-    return nullptr;
-  }
-
-  int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
-  int32_t mIdx = rowIdx;
-
-  // SM103 tile decomposition (128 rows per M-tile, 4 K-positions per K-tile).
-  int32_t mTileIdx = mIdx >> 7;         // mIdx / 128
-  int32_t mLocal = mIdx & 127;          // mIdx % 128
-
-  // SM103 3-level M decomposition: mLocal = m4b + m4a*4 + m8*16
-  int32_t m4b = mLocal & 3;             // mLocal % 4
-  int32_t m4a = (mLocal >> 2) & 3;      // (mLocal / 4) % 4
-  int32_t m8  = (mLocal >> 4) & 7;      // (mLocal / 16) % 8
-
-  int32_t kTileIdx = kIdx >> 2;         // kIdx / 4
-  int32_t innerKIdx = kIdx & 3;         // kIdx % 4
-
-  // Physical offset within the 512-byte tile:
-  //   m8 * 16 + m4a * 128 + m4b * 4 + innerKIdx
-  // Tile base: (mTileIdx * numKTiles + kTileIdx) * 512
-  int64_t SFOffset = (static_cast<int64_t>(mTileIdx) * numKTiles + kTileIdx)
-                         << 9 |
-                     (m8 << 4) | (m4a << 7) | (m4b << 2) | innerKIdx;
-
-  return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
-}
-
 template <class SFType>
 __device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
                                                       int packs_per_row_sf,
@@ -286,15 +286,6 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                "Outer scale stride must be 1 when scales are not transposed");
  }

-  int64_t hidden_size = input.size(-1);
-  TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
-              "hidden_size must be a positive multiple of group_size");
-  int64_t num_tokens = input.numel() / hidden_size;
-  int64_t num_groups = hidden_size / group_size;
-  TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
-              "scales buffer too small: need ", num_tokens * num_groups,
-              " elements, got ", scales.numel());
-
  rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                    var_epsilon, scale_ub, residual,
                                    is_scale_transposed);
@@ -17,11 +17,8 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
                                      int32_t* problem_sizes2,
                                      int32_t* atomic_buffer,
                                      const int topk_length, const int n,
-                                      const int k, const bool is_gated) {
+                                      const int k) {
  int expert_id = blockIdx.x;
-  // For gated activations (gate + up), first GEMM output is 2*n.
-  // For non-gated activations (up only), first GEMM output is n.
-  int const n1 = is_gated ? 2 * n : n;

  int occurrences = 0;
  for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
@@ -34,13 +31,13 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,
    int final_occurrences = atomic_buffer[expert_id];
    if constexpr (!SWAP_AB) {
      problem_sizes1[expert_id * 3] = final_occurrences;
-      problem_sizes1[expert_id * 3 + 1] = n1;
+      problem_sizes1[expert_id * 3 + 1] = 2 * n;
      problem_sizes1[expert_id * 3 + 2] = k;
      problem_sizes2[expert_id * 3] = final_occurrences;
      problem_sizes2[expert_id * 3 + 1] = k;
      problem_sizes2[expert_id * 3 + 2] = n;
    } else {
-      problem_sizes1[expert_id * 3] = n1;
+      problem_sizes1[expert_id * 3] = 2 * n;
      problem_sizes1[expert_id * 3 + 1] = final_occurrences;
      problem_sizes1[expert_id * 3 + 2] = k;
      problem_sizes2[expert_id * 3] = k;
@@ -110,11 +107,13 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
 }

 namespace {
-inline void launch_compute_problem_sizes(
-    const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
-    torch::Tensor& problem_sizes2, torch::Tensor& atomic_buffer,
-    int64_t num_experts, int64_t n, int64_t k, cudaStream_t stream,
-    const bool swap_ab, const bool is_gated) {
+inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,
+                                         torch::Tensor& problem_sizes1,
+                                         torch::Tensor& problem_sizes2,
+                                         torch::Tensor& atomic_buffer,
+                                         int64_t num_experts, int64_t n,
+                                         int64_t k, cudaStream_t stream,
+                                         const bool swap_ab) {
  int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());

  auto const* topk_ptr = topk_ids.data_ptr<int32_t>();
@@ -126,7 +125,7 @@ inline void launch_compute_problem_sizes(
    compute_problem_sizes<SwapAB><<<num_experts, num_threads, 0, stream>>>(
        topk_ptr, ps1_ptr, ps2_ptr, atomic_ptr,
        static_cast<int>(topk_ids.numel()), static_cast<int>(n),
-        static_cast<int>(k), is_gated);
+        static_cast<int>(k));
  });
 }
 }  // namespace
@@ -223,8 +222,7 @@ void get_cutlass_moe_mm_data_caller(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const std::optional<torch::Tensor>& blockscale_offsets) {
  auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
  auto options_int32 =
      torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
@@ -238,7 +236,7 @@ void get_cutlass_moe_mm_data_caller(

  launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,
                               atomic_buffer, num_experts, n, k, stream,
-                               may_swap_ab, is_gated);
+                               may_swap_ab);

  if (blockscale_offsets.has_value()) {
    // fp4 path
@@ -75,8 +75,7 @@ void get_cutlass_moe_mm_data_caller(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const std::optional<torch::Tensor>& blockscale_offsets);

 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
    const torch::Tensor& expert_first_token_offset,
@@ -279,8 +278,7 @@ void get_cutlass_moe_mm_data(
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
    torch::Tensor& input_permutation, torch::Tensor& output_permutation,
    const int64_t num_experts, const int64_t n, const int64_t k,
-    const std::optional<torch::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const std::optional<torch::Tensor>& blockscale_offsets) {
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
@@ -290,7 +288,7 @@ void get_cutlass_moe_mm_data(
  get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
                                 problem_sizes2, input_permutation,
                                 output_permutation, num_experts, n, k,
-                                 blockscale_offsets, is_gated);
+                                 blockscale_offsets);
  return;
 #endif
  TORCH_CHECK_NOT_IMPLEMENTED(
@@ -26,16 +26,6 @@
  #define __HIP__GFX9__
 #endif

-#if defined(__HIPCC__) &&                                                    \
-    (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
-     defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
-  #define __HIP__GFX1X__
-#endif
-
-#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
-  #define __HIP__GFX12__
-#endif
-
 #if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
  #define __HIP__MI3XX__
 #endif
@@ -47,31 +37,15 @@
 #endif

 int get_lds_size() {
-  static const int result = [] {
-    const auto* dprops = at::cuda::getCurrentDeviceProperties();
-    const std::string device_arch = dprops->gcnArchName;
-    return device_arch.find("gfx95") == std::string::npos ? 64 * 1024
-                                                          : 160 * 1024;
-  }();
-  return result;
-}
-
-bool on_gfx1x() {
-  static const bool result = [] {
-    const auto* dprops = at::cuda::getCurrentDeviceProperties();
-    const std::string device_arch = dprops->gcnArchName;
-    return device_arch.find("gfx11") != std::string::npos ||
-           device_arch.find("gfx12") != std::string::npos;
-  }();
-  return result;
-}
-
-bool on_gfx12() {
-  static const bool result = [] {
-    const auto* dprops = at::cuda::getCurrentDeviceProperties();
-    const std::string device_arch = dprops->gcnArchName;
-    return device_arch.find("gfx12") != std::string::npos;
-  }();
+  static bool is_cached = false;
+  static int result;
+  if (is_cached == false) {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    std::string device_arch = dprops->gcnArchName;
+    size_t substring = device_arch.find("gfx95");
+    result = (substring == std::string::npos ? 64 * 1024 : 160 * 1024);
+    is_cached = true;
+  }
  return result;
 }

@@ -312,35 +286,21 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
  return out_c;
 }

-#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__)
-  #define DOT2C(V0, V2, V3)                                          \
-    if constexpr (std::is_same_v<scalar_t, half>) {                  \
-      asm("v_dot2c_f32_f16 %0, %2, %3"                               \
-          : "=v"(V0)                                                 \
-          : "0"(V0), "v"(V2), "v"(V3));                              \
-    } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) { \
-      float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) *  \
-                 __bfloat1622float2(*((__hip_bfloat162*)(&(V3))));   \
-      V0 += (s.x + s.y);                                             \
-    }
-#elif defined(__HIP__GFX1X__)
-  // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12)
-  #define DOT2C(V0, V2, V3)                                               \
-    if constexpr (std::is_same_v<scalar_t, half>) {                       \
-      asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \
-    } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {      \
-      float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) *       \
-                 __bfloat1622float2(*((__hip_bfloat162*)(&(V3))));        \
-      V0 += (s.x + s.y);                                                  \
-    }
-#endif
+#define DOT2C(V0, V2, V3)                                                     \
+  if constexpr (std::is_same_v<scalar_t, half>) {                             \
+    asm("v_dot2c_f32_f16 %0, %2, %3" : "=v"(V0) : "0"(V0), "v"(V2), "v"(V3)); \
+  } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {            \
+    float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) *             \
+               __bfloat1622float2(*((__hip_bfloat162*)(&(V3))));              \
+    V0 += (s.x + s.y);                                                        \
+  }

 // To avoid LLVM silently upcasting to double
 __device__ inline unsigned int min__(uint32_t a, uint32_t b) {
  return min(a, b);
 }

-#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
+#if defined(__HIP__GFX9__)  // TODO: Add NAVI support
 // This version targets cases where A[] fits LDS capacity
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
@@ -482,18 +442,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                                                1);  // row_shr2
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
                                                1);  // row_shr1
-  #if defined(__HIP__GFX9__)
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
                                                1);  // ROW_BCAST15
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
                                                1);  // ROW_BCAST31
-  #else
-          sum[n][y] += __shfl_xor(sum[n][y], 16);
-  #endif
        }
      }

-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -513,10 +469,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        }
      }
    } else {
-  #ifdef __HIP__GFX9__
-    #pragma unroll
+  #pragma unroll
      for (int n = 0; n < N; n++) {
-    #pragma unroll
+  #pragma unroll
        for (int y = 0; y < YTILE; y++) {
          /*float accm1 = 0;
           for (int i=0; i<64; i++)
@@ -543,7 +498,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          sum4[n][y][0] = accm;
        }
      }
-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -558,12 +513,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          }
        }
      }
-  #endif  // __HIP__GFX9__ (MFMA path)
    }
    m += CuCount * _WvPrGrp * YTILE;
  }
 }
-#else
+#else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
 __global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
@@ -574,9 +528,9 @@ __global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap,
                                 const int _WvPrGrp, const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif
+#endif  // defined(__HIP__GFX9__) TODO: Add NAVI support

-#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
+#if defined(__HIP__GFX9__)  // TODO: Add NAVI support
 // This version targets cases where A[] marginally exceeds LDS capacity
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
@@ -703,18 +657,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                                                1);  // row_shr2
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
                                                1);  // row_shr1
-  #if defined(__HIP__GFX9__)
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
                                                1);  // ROW_BCAST15
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
                                                1);  // ROW_BCAST31
-  #else
-          sum[n][y] += __shfl_xor(sum[n][y], 16);
-  #endif
        }
      }

-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -736,10 +686,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        }
      }
    } else {
-  #ifdef __HIP__GFX9__
-    #pragma unroll
+  #pragma unroll
      for (int n = 0; n < N; n++) {
-    #pragma unroll
+  #pragma unroll
        for (int y = 0; y < YTILE; y++) {
          // float accm1 = 0;
          // for (int i=0; i<64; i++)
@@ -764,7 +713,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          sum4[n][y][0] = accm;
        }
      }
-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -781,7 +730,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          }
        }
      }
-  #endif  // __HIP__GFX9__ (MFMA path)
    }

    m += CuCount * _WvPrGrp * YTILE;
@@ -798,7 +746,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  }
 }

-#else
+#else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
 __global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
@@ -808,9 +756,9 @@ __global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap,
                             const int _WvPrGrp, const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif
+#endif  // defined(__HIP__GFX9__) TODO: Add NAVI support

-#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__)
+#if defined(__HIP__GFX9__)  // TODO: Add NAVI support
 // This version targets big A[] cases, where it is much larger than LDS capacity
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
@@ -1056,18 +1004,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                                                1);  // row_shr2
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf,
                                                1);  // row_shr1
-  #if defined(__HIP__GFX9__)
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf,
                                                1);  // ROW_BCAST15
          sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf,
                                                1);  // ROW_BCAST31
-  #else
-          sum[n][y] += __shfl_xor(sum[n][y], 16);
-  #endif
        }
      }

-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -1089,10 +1033,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        }
      }
    } else {
-  #ifdef __HIP__GFX9__
-    #pragma unroll
+  #pragma unroll
      for (int n = 0; n < N; n++) {
-    #pragma unroll
+  #pragma unroll
        for (int y = 0; y < YTILE; y++) {
          float accm = sum4[n][y][0];
          accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf,
@@ -1114,7 +1057,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          sum4[n][y][0] = accm;
        }
      }
-      if (threadIdx.x == (THRDS - 1)) {
+      if (threadIdx.x == 63) {
        scalar_t biases[N][YTILE] = {};
        if (BIAS)
          for (int n = 0; n < N; n++) {
@@ -1131,7 +1074,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          }
        }
      }
-  #endif  // __HIP__GFX9__ (MFMA path)
    }

    m += CuCount * _WvPrGrp * YTILE;
@@ -1148,7 +1090,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    }
  }
 }
-#else
+#else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N>
 __global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
@@ -1159,7 +1101,7 @@ __global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap,
                                 const int _WvPrGrp, const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif
+#endif  // defined(__HIP__GFX9__) TODO: Add NAVI support

 // Find the min val of div2 that doesn't increase N/(div1*div2)
 int mindiv(int N, int div1, int div2) {
@@ -1206,40 +1148,40 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int max_lds_len = get_lds_size() / 2;

-#define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N)                     \
+#define WVSPLITK(_YTILE, _UNRL, _N)                                           \
  {                                                                           \
-    dim3 block(_THRDS, _WVPRGRP);                                             \
-    int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, _WVPRGRP);                 \
+    dim3 block(64, 16);                                                       \
+    int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, 16);                       \
    if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0))               \
-      wvSplitK_hf_sml_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N>        \
+      wvSplitK_hf_sml_<fptype, 64, _YTILE, 16, 8, _UNRL, _N>                  \
          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
                                       By_in, af4, bf4, biasf4, c, __wvPrGrp, \
                                       CuCount);                              \
    else if (Kbp_in * N_in <= max_lds_len * 1.2)                              \
-      wvSplitK_hf_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N>            \
+      wvSplitK_hf_<fptype, 64, _YTILE, 16, 8, _UNRL, _N>                      \
          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
                                       By_in, af4, bf4, biasf4, c, __wvPrGrp, \
                                       CuCount);                              \
    else                                                                      \
-      wvSplitK_hf_big_<fptype, _THRDS, _YTILE, _WVPRGRP, 8, _UNRL, _N>        \
+      wvSplitK_hf_big_<fptype, 64, _YTILE, 16, 8, _UNRL, _N>                  \
          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
                                       By_in, af4, bf4, biasf4, c, __wvPrGrp, \
                                       CuCount);                              \
  }

-#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N)     \
+#define WVSPLIT_TILE(_sYT, __N)                           \
  {                                                       \
    bool fit_lds = (Kbp_in * N_in <= max_lds_len);        \
    if (_sYT <= 1)                                        \
-      WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N)           \
+      WVSPLITK(1, 4, __N)                                 \
    else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \
-      WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N)           \
+      WVSPLITK(2, 2, __N)                                 \
    else if (_sYT <= 4 * 3)                               \
-      WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N)           \
+      WVSPLITK(3, 2, __N)                                 \
    else if (__N == 4)                                    \
-      WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N)           \
+      WVSPLITK(4, 1, __N)                                 \
    else                                                  \
-      WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N)           \
+      WVSPLITK(4, 2, __N)                                 \
  }

  AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] {
@@ -1256,31 +1198,18 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
    // then cut the active waves to balance their distribution...
    int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4);

-    const bool use_wave32 = on_gfx1x();
    switch (N_in) {
      case 1:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 1)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 1)
+        WVSPLIT_TILE(sYT, 1)
        break;
      case 2:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 2)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 2)
+        WVSPLIT_TILE(sYT, 2)
        break;
      case 3:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 3)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 3)
+        WVSPLIT_TILE(sYT, 3)
        break;
      case 4:
-        if (use_wave32)
-          WVSPLIT_TILE_CFG(32, 16, sYT, 4)
-        else
-          WVSPLIT_TILE_CFG(64, 16, sYT, 4)
+        WVSPLIT_TILE(sYT, 4)
        break;
      default:
        throw std::runtime_error(
@@ -1724,7 +1653,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #endif
  }
 }
-#else
+#else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
          int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
 __global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
@@ -1759,8 +1688,6 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
  TORCH_CHECK(in_a.dtype() == torch::kFloat16 ||
              in_a.dtype() == torch::kBFloat16);

-  const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
-
  auto out_c = torch::empty(
      {N_in, M_in},
      torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
@@ -1769,6 +1696,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,

  dim3 grid(CuCount);

+  const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  // const int max_lds_len = get_lds_size() / 2;

@@ -1845,7 +1773,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
  return out_c;
 }

-#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#if defined(__HIP__MI3XX__)  // TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -1889,17 +1817,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

  uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;

+  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
  float sA = *s_A;
  float sB = *s_B;

  while (m < M) {
-  #ifdef __HIP__GFX12__
-    // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
-    float sum[N][YTILE] = {};
-  #else
-    // gfx9: MFMA accumulation
    scalar8 sum[N][YTILE] = {};
-  #endif
    for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
      bigType bigA[N][UNRL] = {};
      bigType bigB[YTILE][UNRL];
@@ -1931,17 +1854,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #pragma unroll
      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
        for (uint32_t n = 0; n < N; n++) {
-  #ifdef __HIP__GFX12__
-          // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
-          for (int y = 0; y < YTILE; ++y) {
-    #pragma unroll
-            for (int i = 0; i < A_CHUNK / 4; i++) {
-              sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
-                  bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
-            }
-          }
-  #else
-          // gfx9: MFMA path
          for (int i = 0; i < A_CHUNK; i += 8) {
            for (int y = 0; y < YTILE; ++y) {
              sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
@@ -1949,33 +1861,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                  0);
            }
          }
-  #endif
        }
      }
    }

    // Final reduction
-  #ifdef __HIP__GFX12__
-    // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
-    for (int n = 0; n < N; n++) {
-      for (int y = 0; y < YTILE; y++) {
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        sum[n][y] += __shfl_xor(sum[n][y], 16);
-      }
-    }
-  #else
-    // gfx9 MFMA reduction
    for (int n = 0; n < N; n++) {
      for (int y = 0; y < YTILE; y++) {
        float accm0 = sum[n][y][0];
@@ -1990,15 +1880,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        sum[n][y][0] = accm0;
      }
    }
-  #endif

-    const bool writeback_lane =
-  #ifdef __HIP__GFX12__
-        threadIdx.x == (THRDS - 1);
-  #else
-        threadIdx.x == 0;
-  #endif
-    if (writeback_lane) {
+    if (threadIdx.x == 0) {
      scalar_t biases[N][YTILE] = {};
      if (BIAS)
        for (int n = 0; n < N; n++) {
@@ -2009,17 +1892,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
      for (int n = 0; n < N; n++) {
        for (int y = 0; y < YTILE; y++) {
          if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __HIP__GFX12__
-          float result = sum[n][y] * sA * sB;
-  #else
-          float result = sum[n][y][0] * sA * sB;
-  #endif
+          sum[n][y][0] *= sA * sB;
          if constexpr (std::is_same_v<scalar_t, half>) {
-            result += __half2float(biases[n][y]);
+            sum[n][y][0] += __half2float(biases[n][y]);
          } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            result += __bfloat162float(biases[n][y]);
+            sum[n][y][0] += __bfloat162float(biases[n][y]);
          }
-          C[m + y + n * M] = __float2s<scalar_t>(result);
+          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
        }
      }
    }
@@ -2027,7 +1906,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    m += CuCount * _WvPrGrp * YTILE;
  }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
+#else   // !defined(__HIP__MI3XX__) TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
@@ -2039,9 +1918,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
                                  const int _WvPrGrp, const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#endif  // defined(__HIP__MI3XX__) TODO: Add NAVI support

-#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#if defined(__HIP__MI3XX__)  // TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
@@ -2084,17 +1963,12 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)

  uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;

+  using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
  float sA = *s_A;
  float sB = *s_B;

  while (m < M) {
-  #ifdef __HIP__GFX12__
-    // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
-    float sum[N][YTILE] = {};
-  #else
-    // gfx9: MFMA accumulation
    scalar8 sum[N][YTILE] = {};
-  #endif
    for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
      bigType bigA[N][UNRL] = {};
      bigType bigB[YTILE][UNRL];
@@ -2128,17 +2002,6 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #pragma unroll
      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
        for (uint32_t n = 0; n < N; n++) {
-  #ifdef __HIP__GFX12__
-          // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
-          for (int y = 0; y < YTILE; ++y) {
-    #pragma unroll
-            for (int i = 0; i < A_CHUNK / 4; i++) {
-              sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8(
-                  bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]);
-            }
-          }
-  #else
-          // gfx9: MFMA path
          for (int i = 0; i < A_CHUNK; i += 8) {
            for (int y = 0; y < YTILE; ++y) {
              sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(
@@ -2146,33 +2009,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
                  0);
            }
          }
-  #endif
        }
      }
    }

    // Final reduction
-  #ifdef __HIP__GFX12__
-    // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
-    for (int n = 0; n < N; n++) {
-      for (int y = 0; y < YTILE; y++) {
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 "
-            : "=v"(sum[n][y])
-            : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y]));
-        sum[n][y] += __shfl_xor(sum[n][y], 16);
-      }
-    }
-  #else
-    // gfx9 MFMA reduction
    for (int n = 0; n < N; n++) {
      for (int y = 0; y < YTILE; y++) {
        float accm0 = sum[n][y][0];
@@ -2187,15 +2028,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        sum[n][y][0] = accm0;
      }
    }
-  #endif

-    const bool writeback_lane =
-  #ifdef __HIP__GFX12__
-        threadIdx.x == (THRDS - 1);
-  #else
-        threadIdx.x == 0;
-  #endif
-    if (writeback_lane) {
+    if (threadIdx.x == 0) {
      scalar_t biases[N][YTILE] = {};
      if (BIAS)
        for (int n = 0; n < N; n++) {
@@ -2206,17 +2040,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
      for (int n = 0; n < N; n++) {
        for (int y = 0; y < YTILE; y++) {
          if (y + m >= M) break;  // To avoid mem access fault.
-  #ifdef __HIP__GFX12__
-          float result = sum[n][y] * sA * sB;
-  #else
-          float result = sum[n][y][0] * sA * sB;
-  #endif
+          sum[n][y][0] *= sA * sB;
          if constexpr (std::is_same_v<scalar_t, half>) {
-            result += __half2float(biases[n][y]);
+            sum[n][y][0] += __half2float(biases[n][y]);
          } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            result += __bfloat162float(biases[n][y]);
+            sum[n][y][0] += __bfloat162float(biases[n][y]);
          }
-          C[m + y + n * M] = __float2s<scalar_t>(result);
+          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
        }
      }
    }
@@ -2224,7 +2054,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    m += CuCount * _WvPrGrp * YTILE;
  }
 }
-#else   // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
+#else   // !defined(__HIP__MI3XX__) TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
          int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
@@ -2236,7 +2066,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
                              const int CuCount) {
  UNREACHABLE_CODE
 }
-#endif  // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
+#endif  // defined(__HIP__MI3XX__) TODO: Add NAVI support

 void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
               const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
@@ -2269,30 +2099,24 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int max_lds_len = get_lds_size();

-#define WVSPLITKQ_IMPL(_THRDS, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
-  {                                                                            \
-    dim3 block(_THRDS, _WvPrGrp);                                              \
-    if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {             \
-      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16));      \
-      wvSplitKQ_hf_sml_<fptype, fp8_t, _THRDS, _YTILEs, _WvPrGrp, 16, _UNRLs,  \
-                        _N><<<grid, block, 0, stream>>>(                       \
-          K_in, Kap_in, Kbp_in, M_in, Bx_in, By_in, b_ptr, a_ptr, bias_ptr,    \
-          c_ptr, s_a, s_b, __wvPrGrp, CuCount);                                \
-    } else {                                                                   \
-      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16));      \
-      wvSplitKQ_hf_<fptype, fp8_t, _THRDS, _YTILEm, _WvPrGrp, 16, _UNRLm, _N>  \
-          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,      \
-                                       By_in, b_ptr, a_ptr, bias_ptr, c_ptr,   \
-                                       s_a, s_b, __wvPrGrp, CuCount);          \
-    }                                                                          \
+#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)             \
+  {                                                                           \
+    dim3 block(64, _WvPrGrp);                                                 \
+    if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {            \
+      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16));     \
+      wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
+          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
+                                       By_in, b_ptr, a_ptr, bias_ptr, c_ptr,  \
+                                       s_a, s_b, __wvPrGrp, CuCount);         \
+    } else {                                                                  \
+      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16));     \
+      wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N>     \
+          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
+                                       By_in, b_ptr, a_ptr, bias_ptr, c_ptr,  \
+                                       s_a, s_b, __wvPrGrp, CuCount);         \
+    }                                                                         \
  }

-#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)      \
-  if (on_gfx12())                                                      \
-    WVSPLITKQ_IMPL(32, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \
-  else                                                                 \
-    WVSPLITKQ_IMPL(64, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)
-
  AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
    using fptype = typename scalar<scalar_t>::type;
    auto c_ptr = reinterpret_cast<fptype*>(out_c.data_ptr());
@@ -2312,10 +2136,10 @@ void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
          WVSPLITKQ(16, 2, 2, 2, 2, 2)
          break;
        case 3:
-          WVSPLITKQ(16, 2, 2, 1, 1, 3)
+          WVSPLITKQ(16, 2, 2, 2, 2, 3)
          break;
        case 4:
-          WVSPLITKQ(16, 2, 2, 1, 1, 4)
+          WVSPLITKQ(16, 2, 2, 2, 2, 4)
          break;
        default:
          throw std::runtime_error(
@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
  // The range of logits within the row.
  int rowStart = 0;
  int seq_len = seqLens[rowIdx / next_n];
-  int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
+  int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;

  // Local pointers to this block
  if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
@@ -0,0 +1,90 @@
+#pragma once
+
+// clang-format will break include orders
+// clang-format off
+#include <cudaTypedefs.h>
+
+#if defined CUDA_VERSION && CUDA_VERSION >= 12020
+#include "sparse_scaled_mm_c3x.cuh"
+
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/transform/device/transform_universal_adapter.hpp"
+#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+
+// clang-format on
+
+using namespace cute;
+using namespace vllm;
+
+using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
+/// Make A structured sparse by replacing elements with 0 and compress it
+template <typename Gemm>
+CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
+  // Checks for conformality
+  TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
+              a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
+  TORCH_CHECK(a.dim() == 2)
+  // Check for strides and alignment
+  TORCH_CHECK(a.stride(0) % 4 == 0)  // Required for semi-structured sparsity
+  TORCH_CHECK(a.stride(1) == 1)
+
+  using GemmKernel = typename Gemm::KernelType;
+  using ElementA = typename Gemm::ElementAB;
+  using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
+
+  int m = a.size(0);
+  int k = a.size(1);
+  using ProblemShape = typename GemmKernel::ProblemShape;
+  ProblemShape prob_shape{m, 1, k, 1};
+
+  int64_t lda = a.stride(0);
+  using StrideA = Stride<int64_t, Int<1>, int64_t>;
+  StrideA a_stride{lda, Int<1>{}, 0};
+
+  using CompressorUtility = typename Gemm::CompressorUtility;
+  CompressorUtility compressor_utility(prob_shape, a_stride);
+
+  // Allocate buffers for the metadata E and the compressed matrix A
+  int ME = compressor_utility.get_metadata_m_physical();
+  int KE = compressor_utility.get_metadata_k_physical();
+  int MC = compressor_utility.get_tensorA_m_physical();
+  int KC = compressor_utility.get_tensorA_k_physical();
+
+  auto const a_meta_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  auto const a_nzs_options =
+      torch::TensorOptions().dtype(a.dtype()).device(a.device());
+
+  auto a_meta = torch::zeros({ME, KE}, a_meta_options);
+  auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
+
+  auto a_ptr = static_cast<ElementA*>(a.data_ptr());
+  auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
+  auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
+
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = a.device().index();
+  hw_info.sm_count =
+      cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
+          hw_info.device_id);
+
+  using Compressor = typename Gemm::Compressor;
+  typename Compressor::Arguments arguments{
+      prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
+
+  Compressor compressor_op;
+  size_t workspace_size = Compressor::get_workspace_size(arguments);
+  auto const workspace_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  auto workspace = torch::empty(workspace_size, workspace_options);
+
+  CUTLASS_CHECK(compressor_op.can_implement(arguments));
+  CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
+  CUTLASS_CHECK(compressor_op.run());
+  CUDA_CHECK(cudaDeviceSynchronize());
+
+  return {a_meta, a_nzs};
+}
+
+#endif
@@ -0,0 +1,307 @@
+// clang-format will break include orders
+// clang-format off
+#include <cudaTypedefs.h>
+
+#if defined CUDA_VERSION && CUDA_VERSION >= 12020
+#include "sparse_scaled_mm_c3x.cuh"
+// clang-format on
+
+using namespace cute;
+using namespace vllm;
+
+struct GemmCallerTraits {
+  using return_type = void;
+
+  template <typename GemmConfig, typename... Args>
+  static return_type invoke(Args&&... args) {
+    return cutlass_sparse_gemm_caller<GemmConfig>(std::forward<Args>(args)...);
+  }
+};
+
+struct GemmCompressorTraits {
+  using return_type = CompressorResult;
+
+  template <typename GemmConfig, typename... Args>
+  static return_type invoke(Args&&... args) {
+    return cutlass_sparse_compress<GemmConfig>(std::forward<Args>(args)...);
+  }
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue,
+          typename DispatchFunc, typename... Args>
+typename DispatchFunc::return_type cutlass_gemm_sm90_fp8_dispatch(
+    uint32_t m, uint32_t n, Args&&... args) {
+  static_assert(std::is_same_v<InType, cutlass::float_e4m3_t>);
+
+  using Cutlass3xGemmDefault =
+      typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM64 =
+      typename sm90_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM128 =
+      typename sm90_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM256 =
+      typename sm90_fp8_config_M256<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM512 =
+      typename sm90_fp8_config_M512<InType, OutType, Epilogue>::Cutlass3xGemm;
+
+  using Cutlass3xGemm1 =
+      typename sm90_fp8_config_1<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm2 =
+      typename sm90_fp8_config_2<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm3 =
+      typename sm90_fp8_config_3<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm4 =
+      typename sm90_fp8_config_4<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm5 =
+      typename sm90_fp8_config_5<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm6 =
+      typename sm90_fp8_config_6<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm7 =
+      typename sm90_fp8_config_7<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemm8 =
+      typename sm90_fp8_config_8<InType, OutType, Epilogue>::Cutlass3xGemm;
+
+  uint32_t const mp2 =
+      std::max(static_cast<uint32_t>(64), next_pow_2(m));  // next power of 2
+
+  if (mp2 <= 64) {
+    if (n == 28672) {
+      return DispatchFunc::template invoke<Cutlass3xGemm2>(
+          std::forward<Args>(args)...);
+    } else if (n == 4096 || n == 6144) {
+      return DispatchFunc::template invoke<Cutlass3xGemm1>(
+          std::forward<Args>(args)...);
+    }
+  } else if (mp2 <= 128) {
+    if (n == 4096) {
+      return DispatchFunc::template invoke<Cutlass3xGemm3>(
+          std::forward<Args>(args)...);
+    } else if (n == 28672) {
+      return DispatchFunc::template invoke<Cutlass3xGemm5>(
+          std::forward<Args>(args)...);
+    } else if (n == 6144) {
+      return DispatchFunc::template invoke<Cutlass3xGemm4>(
+          std::forward<Args>(args)...);
+    }
+  } else if (mp2 <= 256) {
+    if (n == 4096) {
+      return DispatchFunc::template invoke<Cutlass3xGemm6>(
+          std::forward<Args>(args)...);
+    } else if (n == 28672) {
+      return DispatchFunc::template invoke<Cutlass3xGemm8>(
+          std::forward<Args>(args)...);
+    } else if (n == 6144) {
+      return DispatchFunc::template invoke<Cutlass3xGemm7>(
+          std::forward<Args>(args)...);
+    }
+  } else {
+    if (n == 6144 || n == 28672) {
+      return DispatchFunc::template invoke<Cutlass3xGemm8>(
+          std::forward<Args>(args)...);
+    } else if (n == 4096) {
+      return DispatchFunc::template invoke<Cutlass3xGemm7>(
+          std::forward<Args>(args)...);
+    }
+  }
+
+  // Otherwise the default heuristic
+  if (mp2 <= 64) {
+    // n in [1, 64]
+    return DispatchFunc::template invoke<Cutlass3xGemmM64>(
+        std::forward<Args>(args)...);
+  } else if (mp2 <= 128) {
+    // n in (64, 128]
+    return DispatchFunc::template invoke<Cutlass3xGemmM128>(
+        std::forward<Args>(args)...);
+  } else if (mp2 <= 256) {
+    // n in (128, 256]
+    return DispatchFunc::template invoke<Cutlass3xGemmM256>(
+        std::forward<Args>(args)...);
+  } else {
+    // n in (256, inf)
+    return DispatchFunc::template invoke<Cutlass3xGemmM512>(
+        std::forward<Args>(args)...);
+  }
+}
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue,
+          typename DispatchFunc, typename... Args>
+typename DispatchFunc::return_type cutlass_gemm_sm90_16bit_dispatch(
+    uint32_t m, uint32_t n, Args&&... args) {
+  using Cutlass3xGemmDefault =
+      typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
+
+  return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
+      std::forward<Args>(args)...);
+}
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue,
+          typename DispatchFunc, typename... Args>
+typename DispatchFunc::return_type cutlass_gemm_sm90_int8_dispatch(
+    uint32_t m, uint32_t n, Args&&... args) {
+  static_assert(std::is_same_v<InType, int8_t>);
+
+  using Cutlass3xGemmDefault =
+      typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM128 =
+      typename sm90_int8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM64 =
+      typename sm90_int8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM32NBig =
+      typename sm90_int8_config_M32_NBig<InType, OutType,
+                                         Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM32NSmall =
+      typename sm90_int8_config_M32_NSmall<InType, OutType,
+                                           Epilogue>::Cutlass3xGemm;
+
+  bool const is_small_n = n < 8192;
+  uint32_t const mp2 =
+      std::max(static_cast<uint32_t>(32), next_pow_2(m));  // next power of 2
+
+  if (mp2 <= 32) {
+    // m in [1, 32]
+    if (is_small_n) {
+      return DispatchFunc::template invoke<Cutlass3xGemmM32NSmall>(
+          std::forward<Args>(args)...);
+    } else {
+      return DispatchFunc::template invoke<Cutlass3xGemmM32NBig>(
+          std::forward<Args>(args)...);
+    }
+  } else if (mp2 <= 64) {
+    // m in (32, 64]
+    return DispatchFunc::template invoke<Cutlass3xGemmM64>(
+        std::forward<Args>(args)...);
+  } else if (mp2 <= 128) {
+    // m in (64, 128]
+    return DispatchFunc::template invoke<Cutlass3xGemmM128>(
+        std::forward<Args>(args)...);
+  } else {
+    // m in (128, inf)
+    return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
+        std::forward<Args>(args)...);
+  }
+}
+
+// Dispatch to GEMM implementations based on element types
+template <template <typename, typename, typename> typename Epilogue,
+          typename... EpilogueArgs>
+void cutlass_scaled_sparse_mm_sm90_epilogue(torch::Tensor& out,
+                                            torch::Tensor const& a,
+                                            torch::Tensor const& bt_nzs,
+                                            torch::Tensor const& bt_meta,
+                                            EpilogueArgs&&... epilogue_args) {
+  uint32_t const m = out.size(0);
+  uint32_t const n = out.size(1);
+
+  // TODO: add dispatch functions to all of these
+  TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
+  if (a.dtype() == torch::kInt8) {
+    TORCH_CHECK(bt_nzs.dtype() == torch::kInt8);
+
+    if (out.dtype() == torch::kBFloat16) {
+      return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
+                                             Epilogue, GemmCallerTraits>(
+          m, n, out, a, bt_nzs, bt_meta,
+          std::forward<EpilogueArgs>(epilogue_args)...);
+    } else {
+      TORCH_CHECK(out.dtype() == torch::kFloat16);
+      return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::half_t, Epilogue,
+                                             GemmCallerTraits>(
+          m, n, out, a, bt_nzs, bt_meta,
+          std::forward<EpilogueArgs>(epilogue_args)...);
+    }
+  } else if (a.dtype() == torch::kFloat8_e4m3fn) {
+    TORCH_CHECK(bt_nzs.dtype() == torch::kFloat8_e4m3fn);
+
+    if (out.dtype() == torch::kBFloat16) {
+      return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
+                                            cutlass::bfloat16_t, Epilogue,
+                                            GemmCallerTraits>(
+          m, n, out, a, bt_nzs, bt_meta,
+          std::forward<EpilogueArgs>(epilogue_args)...);
+    } else {
+      TORCH_CHECK(out.dtype() == torch::kFloat16);
+      return cutlass_gemm_sm90_fp8_dispatch<
+          cutlass::float_e4m3_t, cutlass::half_t, Epilogue, GemmCallerTraits>(
+          m, n, out, a, bt_nzs, bt_meta,
+          std::forward<EpilogueArgs>(epilogue_args)...);
+    }
+  } else if (a.dtype() == torch::kFloat16) {
+    TORCH_CHECK(bt_nzs.dtype() == torch::kFloat16);
+    TORCH_CHECK(out.dtype() == torch::kFloat16);
+
+    return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
+                                            Epilogue, GemmCallerTraits>(
+        m, n, out, a, bt_nzs, bt_meta,
+        std::forward<EpilogueArgs>(epilogue_args)...);
+  } else {  // a.dtype() == torch::kBFloat16
+    TORCH_CHECK(a.dtype() == torch::kBFloat16);
+    TORCH_CHECK(bt_nzs.dtype() == torch::kBFloat16);
+    TORCH_CHECK(out.dtype() == torch::kBFloat16);
+
+    return cutlass_gemm_sm90_16bit_dispatch<
+        cutlass::bfloat16_t, cutlass::bfloat16_t, Epilogue, GemmCallerTraits>(
+        m, n, out, a, bt_nzs, bt_meta,
+        std::forward<EpilogueArgs>(epilogue_args)...);
+  }
+}
+
+void cutlass_scaled_sparse_mm_sm90(torch::Tensor& out, torch::Tensor const& a,
+                                   torch::Tensor const& bt_nzs,
+                                   torch::Tensor const& bt_meta,
+                                   torch::Tensor const& a_scales,
+                                   torch::Tensor const& b_scales,
+                                   std::optional<torch::Tensor> const& bias) {
+  TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
+  TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
+  TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
+
+  if (bias) {
+    TORCH_CHECK(bias->dtype() == out.dtype(),
+                "CUTLASS scaled_mm bias dtype must match output dtype ",
+                out.dtype());
+    return cutlass_scaled_sparse_mm_sm90_epilogue<
+        c3x::ScaledEpilogueColumnBias>(out, a, bt_nzs, bt_meta, b_scales,
+                                       a_scales, *bias);
+  } else {
+    return cutlass_scaled_sparse_mm_sm90_epilogue<c3x::ScaledEpilogue>(
+        out, a, bt_nzs, bt_meta, b_scales, a_scales);
+  }
+}
+
+CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a) {
+  // These m and n variables are fordispatching to different GEMM algorithms.
+  uint32_t const m = 1;  // Set M to 1 for compression
+  uint32_t const n = a.size(1);
+
+  // Note: For correctness, the compressed format must be invariant in:
+  //  - M, the flattened number of tokens
+  //  - Whether output dtype is fp16 or bf16
+  //  - CUTLASS epilogues
+
+  if (a.dtype() == torch::kInt8) {
+    return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
+                                           c3x::TrivialEpilogue,
+                                           GemmCompressorTraits>(m, n, a);
+  } else if (a.dtype() == torch::kFloat8_e4m3fn) {
+    return cutlass_gemm_sm90_fp8_dispatch<
+        cutlass::float_e4m3_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
+        GemmCompressorTraits>(m, n, a);
+  } else if (a.dtype() == torch::kFloat16) {
+    return cutlass_gemm_sm90_16bit_dispatch<
+        cutlass::bfloat16_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
+        GemmCompressorTraits>(m, n, a);
+  } else {
+    TORCH_CHECK(a.dtype() == torch::kBFloat16,
+                "cutlass_sparse_compress only supports int8, fp8_e4m3, fp16, "
+                "and bf16 datatypes");
+    return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
+                                            c3x::TrivialEpilogue,
+                                            GemmCompressorTraits>(m, n, a);
+  }
+}
+
+#endif
@@ -0,0 +1,570 @@
+#pragma once
+
+// clang-format will break include orders
+// clang-format off
+#include <cudaTypedefs.h>
+
+#include <torch/all.h>
+
+#include <ATen/cuda/CUDAContext.h>
+
+#include "cuda_utils.h"
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+
+#include "cutlass/transform/device/transform_universal_adapter.hpp"
+#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
+
+#include "core/math.hpp"
+#include "cutlass_extensions/cute_utils.cuh"
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
+#include "cutlass_extensions/common.hpp"
+#include "cutlass_extensions/torch_utils.hpp"
+// clang-format on
+
+using namespace cute;
+
+/*
+   This file defines 2:4 sparse GEMM operations using the CUTLASS 3.x API,
+   for NVIDIA GPUs with sm90a (Hopper) or later.
+*/
+
+namespace {
+
+// A wrapper for the GEMM kernel that is used to guard against compilation on
+// architectures that will never use the kernel. The purpose of this is to
+// reduce the size of the compiled binary.
+// __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
+// into code that will be executed on the device where it is defined.
+template <typename Kernel>
+struct enable_sm90_or_later : Kernel {
+  template <typename... Args>
+  CUTLASS_DEVICE void operator()(Args&&... args) {
+#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
+    Kernel::operator()(std::forward<Args>(args)...);
+#endif
+  }
+};
+
+using GemmUniversalMode = cutlass::gemm::GemmUniversalMode;
+
+/*
+ * cutlass_sparse_3x_gemm defines a 2:4 sparse GEMM kernel via CUTLASS
+ * for SM90 Hopper systems.
+ */
+template <typename ElementAB_, typename ElementD_,
+          template <typename, typename, typename> typename Epilogue_,
+          typename TileShape, typename ClusterShape, typename KernelSchedule,
+          typename EpilogueSchedule>
+struct cutlass_sparse_3x_gemm {
+  using ElementAB = ElementAB_;
+  using ElementD = ElementD_;
+  using ElementAcc =
+      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
+                                float>::type;
+
+  using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
+
+  using ElementC = void;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutC_Transpose =
+      typename cutlass::layout::LayoutTranspose<LayoutC>::type;
+
+  using EVTCompute = typename Epilogue::EVTCompute;
+
+  // These are the minimum alignments needed for the kernels to compile
+  static constexpr int AlignmentAB =
+      128 / cutlass::sizeof_bits<ElementAB>::value;
+  static constexpr int AlignmentCD =
+      128 / cutlass::sizeof_bits<ElementD>::value;
+
+  using CollectiveEpilogue =
+      typename cutlass::epilogue::collective::CollectiveBuilder<
+          cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
+          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAcc, float, ElementC, LayoutC_Transpose, AlignmentCD, ElementD,
+          LayoutC_Transpose, AlignmentCD, EpilogueSchedule,
+          EVTCompute>::CollectiveOp;
+
+  static constexpr size_t CEStorageSize =
+      sizeof(typename CollectiveEpilogue::SharedStorage);
+  using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
+      static_cast<int>(CEStorageSize)>;
+
+  // clang-format off
+  using CollectiveMainloop =
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          cutlass::arch::Sm90, cutlass::arch::OpClassSparseTensorOp,
+          ElementAB, cutlass::layout::RowMajor, AlignmentAB,
+          ElementAB, cutlass::layout::ColumnMajor, AlignmentAB,
+          ElementAcc, TileShape, ClusterShape,
+          Stages,
+          KernelSchedule>::CollectiveOp;
+  // clang-format on
+
+  using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
+      cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
+      cutlass::gemm::PersistentScheduler>>;
+
+  struct GemmKernel : public KernelType {};
+
+  // Sparse compressor definitions
+  using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
+  using LayoutTagA = cutlass::layout::RowMajor;
+  using CompressorUtility =
+      cutlass::transform::kernel::StructuredSparseCompressorUtility<
+          typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
+          SparseConfig>;
+  using CompressorKernel =
+      cutlass::transform::kernel::StructuredSparseCompressor<
+          typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
+          SparseConfig, cutlass::arch::Sm90>;
+  using Compressor =
+      cutlass::transform::device::TransformUniversalAdapter<CompressorKernel>;
+};
+
+/*
+ * This class defines kernel to compress a 2:4 sparse matrix.
+ * The particular format is defined by the Gemm template parameter,
+ * which is a cutlass_sparse_3x_gemm.
+ */
+using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
+/// Make A structured sparse by replacing elements with 0 and compress it
+template <typename Gemm>
+CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
+  // Checks for conformality
+  TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
+              a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
+  TORCH_CHECK(a.dim() == 2)
+  // Check for strides and alignment
+  TORCH_CHECK(a.stride(0) % 4 == 0)  // Required for semi-structured sparsity
+  TORCH_CHECK(a.stride(1) == 1)
+
+  using GemmKernel = typename Gemm::KernelType;
+  using ElementA = typename Gemm::ElementAB;
+  using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
+
+  int m = a.size(0);
+  int k = a.size(1);
+  using ProblemShape = typename GemmKernel::ProblemShape;
+  ProblemShape prob_shape{m, 1, k, 1};
+
+  int64_t lda = a.stride(0);
+  using StrideA = Stride<int64_t, Int<1>, int64_t>;
+  StrideA a_stride{lda, Int<1>{}, 0};
+
+  using CompressorUtility = typename Gemm::CompressorUtility;
+  CompressorUtility compressor_utility(prob_shape, a_stride);
+
+  // Allocate buffers for the metadata E and the compressed matrix A
+  int ME = compressor_utility.get_metadata_m_physical();
+  int KE = compressor_utility.get_metadata_k_physical();
+  int MC = compressor_utility.get_tensorA_m_physical();
+  int KC = compressor_utility.get_tensorA_k_physical();
+
+  auto const a_meta_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  auto const a_nzs_options =
+      torch::TensorOptions().dtype(a.dtype()).device(a.device());
+
+  auto a_meta = torch::zeros({ME, KE}, a_meta_options);
+  auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
+
+  auto a_ptr = static_cast<ElementA*>(a.data_ptr());
+  auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
+  auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
+
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = a.device().index();
+  hw_info.sm_count =
+      cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
+          hw_info.device_id);
+
+  using Compressor = typename Gemm::Compressor;
+  typename Compressor::Arguments arguments{
+      prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
+
+  Compressor compressor_op;
+  size_t workspace_size = Compressor::get_workspace_size(arguments);
+  auto const workspace_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  auto workspace = torch::empty(workspace_size, workspace_options);
+
+  CUTLASS_CHECK(compressor_op.can_implement(arguments));
+  CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
+  CUTLASS_CHECK(compressor_op.run());
+  CUDA_CHECK(cudaDeviceSynchronize());
+
+  return {a_meta, a_nzs};
+}
+
+template <typename Gemm, typename... EpilogueArgs>
+void cutlass_sparse_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
+                                torch::Tensor const& bt_nzs,
+                                torch::Tensor const& bt_meta,
+                                EpilogueArgs&&... epilogue_params) {
+  using ElementAB = typename Gemm::ElementAB;
+  using ElementD = typename Gemm::ElementD;
+
+  // Interface stride expected from the argument a (will get transposed)
+  // We compute C^T = B^T * A^T, but we assume B is transposed before
+  // compression and hence the bt_* naming
+  using LayoutB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutA;
+  using LayoutE = typename Gemm::GemmKernel::CollectiveMainloop::LayoutE;
+
+  // M, N, K after transposition
+  int32_t m = out.size(1);
+  int32_t n = out.size(0);
+  int32_t k = a.size(1);
+
+  int64_t lda = a.stride(0);
+  int64_t ldc = out.stride(0);
+
+  using StrideA = Stride<int64_t, Int<1>, int64_t>;
+  using StrideC = Stride<Int<1>, int64_t, int64_t>;
+
+  StrideA a_stride{lda, Int<1>{}, Int<0>{}};
+  StrideC c_stride{Int<1>{}, ldc, Int<0>{}};
+
+  using GemmKernel = typename Gemm::GemmKernel;
+  typename GemmKernel::ProblemShape prob_shape{m, n, k, 1};
+
+  using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
+  using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
+
+  LayoutB b_layout = SparseConfig::fill_layoutA(prob_shape);
+  LayoutE e_layout = SparseConfig::fill_layoutE(prob_shape);
+
+  auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
+  auto b_ptr = static_cast<ElementAB*>(bt_nzs.data_ptr());
+  auto e_ptr = static_cast<ElementE*>(bt_meta.data_ptr());
+  typename GemmKernel::MainloopArguments mainloop_args{
+      b_ptr, b_layout, a_ptr, a_stride, e_ptr, e_layout};
+
+  auto c_ptr = static_cast<ElementD*>(out.data_ptr());
+  typename GemmKernel::EpilogueArguments epilogue_args{
+      Gemm::Epilogue::prepare_args(
+          std::forward<EpilogueArgs>(epilogue_params)...),
+      c_ptr, c_stride, c_ptr, c_stride};
+
+  typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,
+                                      prob_shape, mainloop_args, epilogue_args};
+
+  // Launch the CUTLASS GEMM kernel.
+  using GemmOp = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  GemmOp gemm_op;
+  CUTLASS_CHECK(gemm_op.can_implement(args));
+
+  size_t workspace_size = gemm_op.get_workspace_size(args);
+  auto const workspace_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
+  auto workspace = torch::empty(workspace_size, workspace_options);
+
+  auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
+
+  cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
+  CUTLASS_CHECK(status);
+}
+
+//////////////////////////////////////////////////
+// Gemm Configs are defined below
+//////////////////////////////////////////////////
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_config_default {};
+
+template <typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_config_default<half_t, OutType, Epilogue> {
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_128, _128, _128>;
+  using ClusterShape = Shape<_1, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<half_t, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_config_default<cutlass::bfloat16_t, OutType, Epilogue> {
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_128, _128, _128>;
+  using ClusterShape = Shape<_1, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<cutlass::bfloat16_t, OutType, Epilogue, TileShape,
+                             ClusterShape, KernelSchedule, EpilogueSchedule>;
+};
+
+//////////////////////// Cherry-Picking Kernels ////////////////////////
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_1 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_8, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_2 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_128, _64, _256>;
+  using ClusterShape = Shape<_8, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_3 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_1, _2, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_4 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_64, _128, _256>;
+  using ClusterShape = Shape<_8, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_5 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<_8, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_6 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _128, _256>;
+  using ClusterShape = Shape<_1, _2, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_7 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_8 {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_128, _256, _128>;
+  using ClusterShape = Shape<_8, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+////////////////////////////////////////////////////////////////////////
+
+template <typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_config_default<cutlass::float_e4m3_t, OutType, Epilogue> {
+  // M in (128, inf)
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_128, _128, _128>;
+  using ClusterShape = Shape<_1, _2, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<cutlass::float_e4m3_t, OutType, Epilogue,
+                             TileShape, ClusterShape, KernelSchedule,
+                             EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_M64 {
+  // M in [1, 64]
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_M128 {
+  // M in (64, 128]
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _128, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_M256 {
+  // M in (128, 256]
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_fp8_config_M512 {
+  // M in (256, ]
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule =
+      cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
+  using EpilogueSchedule =
+      typename cutlass::epilogue::TmaWarpSpecializedCooperative;
+  using TileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_config_default<int8_t, OutType, Epilogue> {
+  // For M > 128 and any N
+  using KernelSchedule =
+      typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_128, _128, _128>;
+  using ClusterShape = Shape<_2, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<int8_t, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_int8_config_M128 {
+  // For M in (64, 128] and any N
+  static_assert(std::is_same<InType, int8_t>());
+  using KernelSchedule =
+      typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _128, _128>;
+  using ClusterShape = Shape<_2, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_int8_config_M64 {
+  // For M in (32, 64] and any N
+  static_assert(std::is_same<InType, int8_t>());
+  using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_1, _1, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_int8_config_M32_NBig {
+  // For M in [1, 32] and N >= 8192
+  static_assert(std::is_same<InType, int8_t>());
+  using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _128, _256>;
+  using ClusterShape = Shape<_1, _4, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm90_int8_config_M32_NSmall {
+  // For M in [1, 32] and N < 8192
+  static_assert(std::is_same<InType, int8_t>());
+  using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
+  using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_1, _8, _1>;
+  using Cutlass3xGemm =
+      cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
+                             KernelSchedule, EpilogueSchedule>;
+};
+
+}  // namespace
@@ -0,0 +1,104 @@
+#include <cudaTypedefs.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include "cutlass_extensions/common.hpp"
+
+bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability) {
+  // sparse CUTLASS kernels need exactly hopper and are not forward compatible
+  //   CUDA 12.2 and SM90 (Hopper)
+
+#if defined CUDA_VERSION
+  return CUDA_VERSION >= 12020 && cuda_device_capability == 90;
+#endif
+
+  return false;
+}
+
+#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
+void cutlass_scaled_sparse_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
+                                   torch::Tensor const& b,
+                                   torch::Tensor const& e,
+                                   torch::Tensor const& a_scales,
+                                   torch::Tensor const& b_scales,
+                                   std::optional<torch::Tensor> const& bias);
+
+using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
+CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a);
+#endif
+
+void cutlass_scaled_sparse_mm(torch::Tensor& c, torch::Tensor const& a,
+                              torch::Tensor const& bt_nzs,
+                              torch::Tensor const& bt_meta,
+                              torch::Tensor const& a_scales,
+                              torch::Tensor const& b_scales,
+                              std::optional<torch::Tensor> const& bias) {
+  // Checks for conformality
+  TORCH_CHECK(a.dim() == 2 && bt_nzs.dim() == 2 && c.dim() == 2);
+  TORCH_CHECK(c.size(1) == bt_nzs.size(0) && bt_nzs.size(1) * 2 == a.size(1) &&
+              a.size(0) == c.size(0));
+  TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
+  TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == bt_nzs.size(0));
+
+  // Check for strides and alignment
+  TORCH_CHECK(a.stride(1) == 1 && bt_nzs.stride(1) == 1 &&
+              c.stride(1) == 1);            // Row-major
+  TORCH_CHECK(c.stride(0) % 16 == 0);       // 16 Byte Alignment
+  TORCH_CHECK(bt_nzs.stride(0) % 16 == 0);  // 16 Byte Alignment
+  TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
+
+  if (bias) {
+    TORCH_CHECK(bias->numel() == bt_nzs.size(0) && bias->is_contiguous() &&
+                bias->dim() == 1);
+  }
+
+  at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
+  int32_t version_num = get_sm_version_num();
+
+  // Guard against compilation issues for sm90 kernels
+#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
+  // We build for 9.0a which is not forward compatible, so restrict this to
+  // Hopper only
+  if (version_num == 90) {
+    cutlass_scaled_sparse_mm_sm90(c, a, bt_nzs, bt_meta, a_scales, b_scales,
+                                  bias);
+    return;
+  }
+#endif
+
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "No compiled cutlass_scaled_sparse_mm for a compute capability less than "
+      "CUDA device capability: ",
+      version_num);
+}
+
+std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a) {
+  // Check for strides and alignment
+  TORCH_CHECK(a.stride(1) == 1);      // Row-major
+  TORCH_CHECK(a.stride(0) % 8 == 0);  // 8 Byte Alignment for Compression
+
+  at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
+  int32_t version_num = get_sm_version_num();
+
+  // Guard against compilation issues for sm90 kernels
+#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
+  // We build for 9.0a which is not forward compatible, so restrict this to
+  // Hopper only
+  if (version_num == 90) {
+    std::vector<torch::Tensor> result_tensors;
+
+    auto [a_meta, a_nzs] = cutlass_sparse_compress_sm90(a);
+    result_tensors.push_back(std::move(a_nzs));
+    result_tensors.push_back(std::move(a_meta));
+    return result_tensors;
+  }
+#endif
+
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "No compiled cutlass_sparse_compress for a compute capability equal to "
+      "CUDA device capability: ",
+      version_num);
+}
@@ -6,41 +6,6 @@
 #include <torch/library.h>
 #include <torch/version.h>

-// Forward declarations for per-SM NVFP4 GEMM and quantization entry points.
-// Defined in nvfp4_scaled_mm_kernels.cu / nvfp4_quant_entry.cu and only
-// compiled when ENABLE_NVFP4_SM100 is set.
-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha);
-void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
-                                  torch::Tensor const& B,
-                                  torch::Tensor const& A_sf,
-                                  torch::Tensor const& B_sf,
-                                  torch::Tensor const& alpha);
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf);
-void scaled_fp4_quant_sm103a_out(torch::Tensor const& input,
-                                 torch::Tensor const& input_sf,
-                                 torch::Tensor& output,
-                                 torch::Tensor& output_sf);
-// PDL-enabled variants (ProgrammaticStreamSerialization).
-void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D,
-                                      torch::Tensor const& A,
-                                      torch::Tensor const& B,
-                                      torch::Tensor const& A_sf,
-                                      torch::Tensor const& B_sf,
-                                      torch::Tensor const& alpha);
-std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_pdl_func(
-    torch::Tensor const& input, torch::Tensor const& input_sf);
-void scaled_fp4_quant_sm103a_pdl_out(torch::Tensor const& input,
-                                     torch::Tensor const& input_sf,
-                                     torch::Tensor& output,
-                                     torch::Tensor& output_sf);
-#endif
-
 // Note on op signatures:
 // The X_meta signatures are for the meta functions corresponding to op X.
 // They must be kept in sync with the signature for X. Generally, only
@@ -338,6 +303,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      ") -> Tensor");
  // conditionally compiled so impl registration is in source file

+  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
+  ops.impl("permute_cols", torch::kCUDA, &permute_cols);
+
  // Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
  ops.def(
      "marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
@@ -451,65 +419,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                      Tensor alpha) -> ()");
  ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);

-#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
-  // SM100-specific entry point (B200 / Blackwell, SM100 SF layout)
-  ops.def(
-      "cutlass_scaled_fp4_mm_sm100a(Tensor! out, Tensor a, Tensor b,"
-      "                             Tensor block_scale_a, Tensor block_scale_b,"
-      "                             Tensor alpha) -> ()");
-  ops.impl("cutlass_scaled_fp4_mm_sm100a", torch::kCUDA,
-           &cutlass_scaled_fp4_mm_sm100a);
-
-  // SM103-specific entry point (B300 / Blackwell Ultra, SM103 SF layout)
-  ops.def(
-      "cutlass_scaled_fp4_mm_sm103a(Tensor! out, Tensor a, Tensor b,"
-      "                             Tensor block_scale_a, Tensor block_scale_b,"
-      "                             Tensor alpha) -> ()");
-  ops.impl("cutlass_scaled_fp4_mm_sm103a", torch::kCUDA,
-           &cutlass_scaled_fp4_mm_sm103a);
-
-  // SM103-native quantization: produces SM103-layout scale factors directly.
-  ops.def(
-      "scaled_fp4_quant_sm103(Tensor input,"
-      "                       Tensor input_scale) -> (Tensor, Tensor)");
-  ops.impl("scaled_fp4_quant_sm103", torch::kCUDA, &scaled_fp4_quant_sm103a_func);
-
-  ops.def(
-      "scaled_fp4_quant_sm103.out(Tensor input,"
-      "                           Tensor input_scale,"
-      "                           *, Tensor(a!) output, Tensor(b!) output_scale)"
-      " -> ()");
-  ops.impl("scaled_fp4_quant_sm103.out", torch::kCUDA,
-           &scaled_fp4_quant_sm103a_out);
-
-  // PDL-enabled SM103 GEMM: launched with ProgrammaticStreamSerialization
-  // so the next kernel on the stream can overlap with this GEMM's tail.
-  ops.def(
-      "cutlass_scaled_fp4_mm_sm103a_pdl(Tensor! out, Tensor a, Tensor b,"
-      "                                 Tensor block_scale_a,"
-      "                                 Tensor block_scale_b,"
-      "                                 Tensor alpha) -> ()");
-  ops.impl("cutlass_scaled_fp4_mm_sm103a_pdl", torch::kCUDA,
-           &cutlass_scaled_fp4_mm_sm103a_pdl);
-
-  // PDL-enabled SM103 quantization: launched with
-  // ProgrammaticStreamSerialization so the subsequent GEMM can begin
-  // before this quant kernel completes.
-  ops.def(
-      "scaled_fp4_quant_sm103_pdl(Tensor input,"
-      "                           Tensor input_scale) -> (Tensor, Tensor)");
-  ops.impl("scaled_fp4_quant_sm103_pdl", torch::kCUDA,
-           &scaled_fp4_quant_sm103a_pdl_func);
-
-  ops.def(
-      "scaled_fp4_quant_sm103_pdl.out(Tensor input,"
-      "                               Tensor input_scale,"
-      "                               *, Tensor(a!) output,"
-      "                               Tensor(b!) output_scale) -> ()");
-  ops.impl("scaled_fp4_quant_sm103_pdl.out", torch::kCUDA,
-           &scaled_fp4_quant_sm103a_pdl_out);
-#endif
-
  // cutlass nvfp4 block scaled group GEMM
  ops.def(
      "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
@@ -580,8 +489,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                        Tensor! problem_sizes1, Tensor! problem_sizes2, "
      "                        Tensor! input_permutation, "
      "                        Tensor! output_permutation, int num_experts, "
-      "                        int n, int k, Tensor? blockscale_offsets, "
-      "                        bool is_gated) -> ()");
+      "                        int n, int k, Tensor? blockscale_offsets) -> "
+      "()");
  ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

  // compute per-expert problem sizes from expert_first_token_offset
@@ -617,6 +526,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("cutlass_scaled_mm_supports_block_fp8",
           &cutlass_scaled_mm_supports_block_fp8);

+  // Check if cutlass sparse scaled_mm is supported for CUDA devices of the
+  // given capability
+  ops.def(
+      "cutlass_sparse_scaled_mm_supported(int cuda_device_capability) -> bool");
+  ops.impl("cutlass_sparse_scaled_mm_supported",
+           &cutlass_sparse_scaled_mm_supported);
+
+  // CUTLASS sparse GEMM, supporting symmetric per-tensor or per-row/column
+  // quantization, as well as bias
+  ops.def(
+      "cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
+      "                         Tensor bt_nzs,"
+      "                         Tensor bt_meta, Tensor a_scales,"
+      "                         Tensor b_scales, Tensor? bias) -> ()");
+  ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
+
+  // CUTLASS sparse matrix compressor
+  ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
+  ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
+
  // SM100 CUTLASS MLA decode
  ops.def(
      "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
@@ -635,21 +564,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

  // Compute NVFP4 block quantized tensor.
  ops.def(
-      "scaled_fp4_quant(Tensor input,"
-      "                 Tensor input_scale, bool "
-      "is_sf_swizzled_layout) -> (Tensor, Tensor)");
-  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
-
-  // Out variant
-  // TODO: Add {at::Tag::out_variant} tag and update all call sites
-  // to use the functional variant once vLLM upgrades PyTorch.
-  // See pytorch/pytorch#176117.
-  ops.def(
-      "scaled_fp4_quant.out(Tensor input,"
-      "                     Tensor input_scale, bool "
-      "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
-      "-> ()");
-  ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
+      "scaled_fp4_quant(Tensor! output, Tensor input,"
+      "                 Tensor! output_scale, Tensor input_scale, bool "
+      "is_sf_swizzled_layout) -> ()");
+  ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);

  // Compute NVFP4 experts quantization.
  ops.def(
@@ -667,16 +585,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
           &silu_and_mul_scaled_fp4_experts_quant);

-  // SM100 <-> SM103 scale factor layout conversion (B300 / Blackwell Ultra)
-  ops.def(
-      "convert_sf_layout_sm100_to_sm103(Tensor(a!) dst, Tensor src) -> ()");
-  ops.impl("convert_sf_layout_sm100_to_sm103", torch::kCUDA,
-           &convert_sf_layout_sm100_to_sm103);
-  ops.def(
-      "convert_sf_layout_sm103_to_sm100(Tensor(a!) dst, Tensor src) -> ()");
-  ops.impl("convert_sf_layout_sm103_to_sm100", torch::kCUDA,
-           &convert_sf_layout_sm103_to_sm100);
-
  // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
  // of the given capability
  ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
@@ -620,7 +620,7 @@ RUN set -eux; \
 ARG BITSANDBYTES_VERSION_X86=0.46.1
 ARG BITSANDBYTES_VERSION_ARM64=0.42.0
 ARG TIMM_VERSION=">=1.0.17"
-ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
+ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
 RUN --mount=type=cache,target=/root/.cache/uv \
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
@@ -628,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
    fi; \
    uv pip install --system accelerate hf_transfer modelscope \
-        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
+        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"

 # ============================================================
 # VLLM INSTALLATION (depends on build stage)
@@ -9,7 +9,6 @@
 #
 # Build targets:
 #   vllm-openai (default): used for serving deployment
-#   vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
 #   vllm-test: used for CI tests
 #   vllm-dev: used for development
 #
@@ -223,19 +222,3 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
 LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"

 ENTRYPOINT ["vllm", "serve"]
-
-
-######################### ZEN CPU PYPI IMAGE #########################
-FROM vllm-openai AS vllm-openai-zen
-
-ARG TARGETARCH
-
-RUN if [ "$TARGETARCH" != "amd64" ]; then \
-        echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
-        exit 1; \
-    fi
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install "vllm[zen]"
-
-ENTRYPOINT ["vllm", "serve"]
@@ -44,7 +44,7 @@ ENV DEBIAN_FRONTEND=noninteractive

 # Install Python and other dependencies
 RUN apt-get update -y \
-    && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \
+    && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev \
    && for i in 1 2 3; do \
        add-apt-repository -y ppa:deadsnakes/ppa && break || \
        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
@@ -76,22 +76,19 @@ ENV UV_LINK_MODE="copy"
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \
    --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \
-    --mount=type=bind,src=requirements/xpu-test.in,target=/workspace/vllm/requirements/xpu-test.in \
    uv pip install --upgrade pip && \
-    uv pip install -r requirements/xpu.txt && \
-    uv pip compile /workspace/vllm/requirements/xpu-test.in \
-        -o /workspace/vllm/requirements/xpu-test.txt \
-        -c /workspace/vllm/requirements/xpu.txt \
-        --index-strategy unsafe-best-match \
-        --extra-index-url ${PIP_EXTRA_INDEX_URL} \
-        --python-version ${PYTHON_VERSION} && \
-    uv pip install grpcio-tools protobuf nanobind && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
    source /opt/intel/oneapi/setvars.sh --force && \
    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
-    export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
-    uv pip install --no-build-isolation -r /workspace/vllm/requirements/xpu-test.txt
-
-
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1

 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"

@@ -83,7 +83,7 @@
      "default": ">=1.0.17"
    },
    "RUNAI_MODEL_STREAMER_VERSION": {
-      "default": ">=0.15.7"
+      "default": ">=0.15.3"
    }
  }
 }
@@ -25,7 +25,7 @@ nav:
    - Models:
      - models/supported_models.md
      - models/generative_models.md
-      - Pooling Models: models/pooling_models
+      - models/pooling_models.md
      - models/extensions
      - Hardware Supported Models:
        - models/hardware_supported_models/*
@@ -75,7 +75,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment
 vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:

 ```bash
-uv pip install pre-commit>=4.5.1
+uv pip install pre-commit
 pre-commit install
 ```

@@ -37,7 +37,7 @@ For [generative models](../../models/generative_models.md), there are two levels

 #### Pooling models

-For [pooling models](../../models/pooling_models/README.md), we simply check the cosine similarity, as defined in [tests/models/utils.py](../../../tests/models/utils.py).
+For [pooling models](../../models/pooling_models.md), we simply check the cosine similarity, as defined in [tests/models/utils.py](../../../tests/models/utils.py).

 ### Multi-modal processing

@@ -3,10 +3,6 @@
 !!! warning
    Profiling is only intended for vLLM developers and maintainers to understand the proportion of time spent in different parts of the codebase. **vLLM end-users should never turn on profiling** as it will significantly slow down the inference.

-!!! tip "Choosing a profiler"
-    - Use **Nsight Systems** for low-overhead, performance-critical profiling.
-    - Use **PyTorch Profiler** for medium-overhead profiling with richer debugging information (e.g., stack traces, memory, shapes). Note that enabling these features adds overhead and is not recommended for benchmarking.
-
 ## Profile with PyTorch Profiler

 We support tracing vLLM workers using different profilers. You can enable profiling by setting the `--profiler-config` flag when launching the server.
@@ -127,8 +127,8 @@ Priority is **1 = highest** (tried first).
 | 3 | `FLASH_ATTN_MLA` |
 | 4 | `FLASHMLA` |
 | 5 | `TRITON_MLA` |
-| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
-| 7 | `FLASHMLA_SPARSE` |
+| 6 | `FLASHMLA_SPARSE` |
+| 7 | `FLASHINFER_MLA_SPARSE` |

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

@@ -140,8 +140,6 @@ Priority is **1 = highest** (tried first).
 | 4 | `TRITON_MLA` |
 | 5 | `FLASHMLA_SPARSE` |

-> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
->
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.

 ## Legend
@@ -166,18 +164,18 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | 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` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
-| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
+| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | 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 | ✅ | ✅ | ❌ | All | 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` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
+| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |

 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -206,14 +204,14 @@ configuration.

 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
-| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
-| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
-| `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` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `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 |
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
+| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
@@ -51,8 +51,11 @@ For example:
 **1. Attention:**

 ```python
+--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
+
 --8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"

+--8<-- "vllm/model_executor/models/deepencoder.py:rel_pos_attention"
 ```

 **2. Activation:**
@@ -167,16 +170,6 @@ For example:
 --8<-- "vllm/model_executor/layers/rotary_embedding/common.py:apply_rotary_emb"
 ```

-**12. Encoder:**
-
-```python
--8<-- "vllm/model_executor/models/deepencoder2.py:qwen2_decoder"
-
--8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
-
--8<-- "vllm/model_executor/models/deepencoder.py:rel_pos_attention"
-```
-
 ## Guidelines for Implementing a New CustomOp

 ### Implement a New CustomOp in vLLM
@@ -244,7 +244,6 @@ statistics relating to that iteration:
  prefill in this iteration. However, we calculate this interval
  relative to when the request was first received by the frontend
  (`arrival_time`) in order to account for input processing time.
-  Currently `arrival_time` starts when tokenization begins.

 For any requests that were completed in a given iteration, we also
 record:
@@ -588,7 +587,7 @@ see:
 - [Benchmarking LLM Workloads for Performance Evaluation and Autoscaling in Kubernetes](https://docs.google.com/document/d/1k4Q4X14hW4vftElIuYGDu5KDe2LtV1XammoG-Xi3bbQ)
 - [Inference Perf](https://github.com/kubernetes-sigs/wg-serving/tree/main/proposals/013-inference-perf)
 - <https://github.com/vllm-project/vllm/issues/5041> and <https://github.com/vllm-project/vllm/pull/12726>.
-
+  
 This is a non-trivial topic. Consider this comment from Rob:

 > I think this metric should focus on trying to estimate what the max
@@ -33,10 +33,9 @@ th {
 | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
 | ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- |
 | naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
-| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
-| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
-| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
-| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
+| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
+| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
+| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |

 !!! info "Table key"
    1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -88,8 +87,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
 | gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
-| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
-| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
+| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
+| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
 | naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |

@@ -103,7 +102,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k

 ## Modular Kernel "families"

-The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts.
+The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts.

 | backend | `FusedMoEPrepareAndFinalizeModular` subclasses | `FusedMoEExpertsModular` subclasses |
 | ------- | ---------------------------------------------- | ----------------------------------- |
@@ -29,9 +29,13 @@ To compile a multimodal component such as an encoder, we follow the same mechani
 1. The `@support_torch_compile` decorator should include `enable_if=should_torch_compile_mm_encoder`. This will gate the compilation behind our
 `compile_mm_encoder` configuration

-2. The `@support_torch_compile` decorator should include `is_encoder=True` for encoder components. This is needed for compile range integration
-(see Compile Range Integration). The decorator automatically uses the class name as the cache directory prefix, avoiding collisions between
-independently compiled sub-modules (e.g. vision encoder components vs the text backbone).
+2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
+relies on caching artifacts to reduce start time, we must properly propagate the `<component_name>` information to the cache in order to avoid collisions
+with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
+components (see Compile Range Integration).
+
+3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
+for torch.compile integration.

 ### CompilationConfig

@@ -56,8 +60,8 @@ tradeoff
 ### Compile ranges

 The torch.compile integration will try to rely on max_batch_size to infer compilation ranges for dynamic shapes; however, for modules used in the encoder, this
-shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the
-`@support_torch_compile` decorator to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).
+shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the `set_model_tag`
+to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).

 !!! note
    We may seek to tighten this range for better performance in the future
@@ -36,14 +36,14 @@ th:not(:first-child) {
 }
 </style>

-| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models/README.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
+| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
 | - | - | - | - | - | - | - | - | - | - | - | - | - | - | - | - |
 | [CP](../configuration/optimization.md#chunked-prefill) | ✅ | | | | | | | | | | | | | | |
 | [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
 | [SD](speculative_decoding/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
 | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
-| [pooling](../models/pooling_models/README.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
+| [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ✅ | ✅ | ✅ | | | | | | | | |
 | <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | | |
 | <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | | | | | | |
@@ -66,7 +66,7 @@ th:not(:first-child) {
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
 | CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
-| [pooling](../models/pooling_models/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
 | [mm](multimodal_inputs.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ |
@@ -389,17 +389,3 @@ vllm serve model --enable-lora --max-lora-rank 64
 # Bad: unnecessarily high, wastes memory
 vllm serve model --enable-lora --max-lora-rank 256
 ```
-
-### Restricting LoRA to Specific Modules
-
-The `--lora-target-modules` parameter allows you to restrict which model modules have LoRA applied at deployment time. This is useful for performance tuning when you only need LoRA on specific layers:
-
-```bash
-# Apply LoRA only to output projection layers
-vllm serve model --enable-lora --lora-target-modules o_proj
-
-# Apply LoRA to multiple specific modules
-vllm serve model --enable-lora --lora-target-modules o_proj qkv_proj down_proj
-```
-
-When `--lora-target-modules` is not specified, LoRA will be applied to all supported modules in the model. This parameter accepts module suffixes (the last component of the module name), such as `o_proj`, `qkv_proj`, `gate_proj`, etc.
@@ -137,3 +137,6 @@ llm = LLM("facebook/opt-125m", quantization="fp8")
 result = llm.generate("Hello, my name is")
 print(result[0].outputs[0].text)
 ```
+
+!!! warning
+    Currently, we load the model at original precision before quantizing down to 8-bits, so you need enough memory to load the whole model.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Nick HillandGitHub	95ef052724	Merge branch 'main' into mrv2-ci-test	2026-03-14 10:29:16 -07:00
Nick Hill	47bbcbd7ef	[WIP] CI test run with Model Runner V2 Expected to fail many tests Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-02-27 14:48:27 -08:00