Bump actions/setup-python from 6.1.0 to 6.2.0

Bumps [actions/setup-python](https://github.com/actions/setup-python) from 6.1.0 to 6.2.0. - [Release notes](https://github.com/actions/setup-python/releases) - [Commits](https://github.com/actions/setup-python/compare/83679a892e2d95755f2dac6acb0bfd1e9ac5d548...a309ff8b426b58ec0e2a45f0f869d46889d02405) --- updated-dependencies: - dependency-name: actions/setup-python dependency-version: 6.2.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com>
2026-06-02 18:31:03 +00:00
928 changed files with 15418 additions and 46328 deletions
@@ -1,23 +0,0 @@
-name: vllm_rocm_ci
-job_dirs:
-  - ".buildkite/hardware_tests"
-run_all_patterns:
-  - "docker/Dockerfile.rocm"
-  - "docker/Dockerfile.rocm_base"
-  - "docker/ci-rocm.hcl"
-  - "docker/docker-bake-rocm.hcl"
-  - ".buildkite/hardware_tests/amd.yaml"
-  - ".buildkite/scripts/ci-bake-rocm.sh"
-  - ".buildkite/scripts/hardware_ci/run-amd-test.py"
-  - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
-  - "CMakeLists.txt"
-  - "requirements/common.txt"
-  - "requirements/rocm.txt"
-  - "requirements/build/rocm.txt"
-  - "requirements/test/rocm.txt"
-  - "setup.py"
-  - "csrc/"
-  - "cmake/"
-run_all_exclude_patterns:
-  - "csrc/cpu/"
-  - "cmake/cpu_extension.cmake"
@@ -1,73 +1,42 @@
-group: Hardware - AMD Build
+group: Hardware - AMD Build 
 steps:
-  # Ensure ci_base is up-to-date before building the test image.
-  # Compares a content hash of ci_base-affecting files against the remote
-  # image label. If hashes match the build is skipped (< 30 s); if they
-  # differ ci_base is rebuilt and pushed automatically.
-  - label: "AMD: :docker: ensure ci_base"
-    key: ensure-ci-base-amd
+  - label: "AMD: :docker: build image"
+    key: image-build-amd
    depends_on: []
    device: amd_cpu
    no_plugin: true
    commands:
-      - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
-    env:
-      DOCKER_BUILDKIT: "1"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      REMOTE_VLLM: "1"
-      VLLM_BRANCH: "$BUILDKITE_COMMIT"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1
-
-  - label: "AMD: :docker: build test image and artifacts"
-    key: image-build-amd
-    depends_on:
-      - ensure-ci-base-amd
-    device: amd_cpu
-    no_plugin: true
-    commands:
-      - |
-        if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
-          echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
-          IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
-        else
-          bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
+    - >
+      docker build
+      --build-arg max_jobs=16
+      --build-arg REMOTE_VLLM=1
+      --build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942;gfx950'
+      --build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
+      --tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
+      -f docker/Dockerfile.rocm
+      --target test
+      --no-cache
+      --progress plain .
+    - |
+      docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
+        if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
+        if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
+        command -v python3
+        command -v uv
+        command -v pytest
+        if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
+          echo No ROCm CLI found in image >&2
+          exit 1
        fi
-      - |
-        docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
-          if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
-          if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
-          if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
-          command -v python3
-          command -v uv
-          command -v pytest
-          if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
-            echo No ROCm CLI found in image >&2
-            exit 1
-          fi
-          python3 - <<PY
-        import torch, vllm
-        print(torch.__version__)
-        print(vllm.__version__)
-        PY
-          echo AMD image smoke OK
-        '
+        python3 - <<PY
+      import torch, vllm
+      print(torch.__version__)
+      print(vllm.__version__)
+      PY
+        echo AMD image smoke OK
+      '
+    - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
    env:
      DOCKER_BUILDKIT: "1"
-      VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
-      PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
-      IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
-      REMOTE_VLLM: "1"
-      VLLM_BRANCH: "$BUILDKITE_COMMIT"
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 1
-        - exit_status: -10  # Agent was lost
-          limit: 1
@@ -28,19 +28,18 @@ steps:
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"

-# Note: SDE can't be downloaded from CI host because of AWS WAF
-# - label: CPU-Compatibility Tests
-#   depends_on: []
-#   device: intel_cpu
-#   no_plugin: true
-#   source_file_dependencies:
-#   - cmake/cpu_extension.cmake
-#   - setup.py
-#   - vllm/platforms/cpu.py
-#   commands:
-#     - |
-#       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
-#       bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
+- label: CPU-Compatibility Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - cmake/cpu_extension.cmake
+  - setup.py
+  - vllm/platforms/cpu.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"

 - label: CPU-Language Generation and Pooling Model Tests
  depends_on: []
@@ -1,22 +0,0 @@
-group: Basic Correctness
-depends_on:
-  - image-build-xpu
-steps:
- label: XPU Sleep Mode
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/basic_correctness/test_cumem.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
@@ -1,23 +0,0 @@
-group: Expert Parallelism
-depends_on: 
-  - image-build-xpu
-steps:
- label: EPLB Algorithm
-  key: eplb-algorithm
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_algo.py
-  - tests/distributed/test_eplb_utils.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'cd tests &&
-      pytest -v -s distributed/test_eplb_algo.py'
@@ -38,17 +38,7 @@ steps:
    REPO: "vllm-ci-test-repo"
    VLLM_TEST_DEVICE: "xpu"
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/engine/
-    - vllm/inputs/
-    - vllm/logger.py
-    - vllm/model_executor/
-    - vllm/platforms/
-    - vllm/sampling_params.py
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - tests/v1/sample
    - tests/v1/logits_processors
    - tests/v1/test_oracle.py
@@ -62,126 +52,4 @@ steps:
      pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
      pytest -v -s v1/test_oracle.py &&
      pytest -v -s v1/test_request.py &&
-      pytest -v -s v1/test_outputs.py &&
-      pytest -v -s v1/sample/test_topk_topp_sampler.py'
-
- label: XPU CPU Offload
-  timeout_in_minutes: 60
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-    - vllm/
-    - vllm/v1/kv_offload/
-    - vllm/v1/kv_connector/
-    - tests/v1/kv_offload/
-    - tests/v1/kv_connector/unit/test_offloading_connector.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
-      cd tests &&
-      pytest -v -s v1/kv_offload &&
-      pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
-
- label: Regression
-  key: regression
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - tests/test_regression
-  commands:
-  - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install modelscope &&
-      cd tests &&
-      pytest -v -s test_regression.py'
-
- label: Metrics, Tracing (2 GPUs)
-  key: metrics-tracing-2-gpus
-  timeout_in_minutes: 30
-  num_devices: 2
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/tracing/
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - tests/v1/tracing
-  commands:
-  - >-
-    bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-    'pip install opentelemetry-sdk\>=1.26.0 opentelemetry-api\>=1.26.0 opentelemetry-exporter-otlp\>=1.26.0 opentelemetry-semantic-conventions-ai\>=0.4.1 &&
-    cd tests &&
-    pytest -v -s v1/tracing'
-
- label: Async Engine, Inputs, Utils, Worker
-  key: async-engine-inputs-utils-worker
-  timeout_in_minutes: 30
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/assets/
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/tokenizers/
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - tests/detokenizer
-  - tests/multimodal
-  - tests/utils_
-  commands:
-  - >-
-    bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-    'cd tests &&
-    pip install av &&
-    pytest -v -s detokenizer &&
-    pytest -v -s -m "not cpu_test" ./multimodal &&
-    pytest -v -s utils_ --ignore=utils_/test_mem_utils.py'
+      pytest -v -s v1/test_outputs.py'
@@ -1,111 +0,0 @@
-group: Models - Multimodal
-depends_on: 
-  - image-build-xpu
-steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
-  key: multi-modal-models-standard-1-qwen2
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
-      pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen2" &&
-      pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model'
-
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
-  key: multi-modal-models-standard-2-qwen3-gemma
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
-      pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model'
-
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
-  key: multi-modal-models-standard-3-llava-qwen2-vl
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
-      pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" &&
-      pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model'
-
- label: "Multi-Modal Models (Standard) 4: other + whisper"
-  key: multi-modal-models-standard-4-other-whisper
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      cd tests &&
-      pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing'
-
- label: Multi-Modal Processor # 44min
-  key: multi-modal-processor
-  timeout_in_minutes: 45
-  device: intel_gpu
-  no_plugin: true
-  working_dir: "."
-  env:
-    REGISTRY: "public.ecr.aws/q9t5s3a7"
-    REPO: "vllm-ci-test-repo"
-    VLLM_TEST_DEVICE: "xpu"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  - tests/models/registry.py
-  commands:
-    - >-
-      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'pip install av matplotlib ftfy git+https://github.com/TIGER-AI-Lab/Mantis.git &&
-      pip install open-clip-torch --no-deps &&
-      cd tests &&
-      pytest -v -s models/multimodal/processing/test_tensor_schema.py
-      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4]"
-      --deselect "tests/models/multimodal/processing/test_tensor_schema.py::test_model_tensor_schema[Qwen/Qwen2.5-Omni-7B-AWQ]"
-      --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB'
-  parallelism: 4
@@ -40,9 +40,7 @@ steps:
        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
-        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
-        VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
-        python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
+        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
        '
  - label: "XPU V1 test"
    depends_on:
@@ -85,5 +83,5 @@ steps:
        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
        'pip install av &&
        cd tests &&
-        pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
+        pytest -v -s entrypoints/openai/chat_completion/test_audio_in_video.py &&
        pytest -v -s benchmarks/test_serve_cli.py'
@@ -2,6 +2,7 @@
    {
        "test_name": "latency_llama8B_tp1",
        "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
            "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
            "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
            "VLLM_CPU_KVCACHE_SPACE": 40
@@ -2,6 +2,7 @@
    {
        "test_name": "latency_llama8B_tp2",
        "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
 	    "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
 	    "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
 	    "VLLM_CPU_SGL_KERNEL": 1,
@@ -13,6 +13,7 @@
      200
    ],
    "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
      "VLLM_CPU_SGL_KERNEL": 1,
@@ -5,6 +5,7 @@
    ],
    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
    "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120
    },
    "server_parameters": {
@@ -9,6 +9,7 @@
      128
    ],
    "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
      "VLLM_CPU_SGL_KERNEL": 1,
@@ -5,6 +5,7 @@
    ],
    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
    "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
      "VLLM_CPU_SGL_KERNEL": 1,
@@ -5,6 +5,7 @@
    ],
    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
    "server_environment_variables": {
+      "VLLM_RPC_TIMEOUT": 100000,
      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
      "VLLM_CPU_SGL_KERNEL": 1,
@@ -2,6 +2,7 @@
    {
        "test_name": "throughput_llama8B_tp1",
        "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
            "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
            "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
            "VLLM_CPU_KVCACHE_SPACE": 40
@@ -2,6 +2,7 @@
    {
        "test_name": "throughput_llama8B_tp2",
        "environment_variables": {
+            "VLLM_RPC_TIMEOUT": 100000,
 	    "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
 	    "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
 	    "VLLM_CPU_SGL_KERNEL": 1,
@@ -28,10 +28,8 @@
 ###############################################################################
 set -o pipefail

-# Export Python path for commands that run directly on the host. Containerized
-# tests set this to /vllm-workspace below so spawned Python processes do not
-# depend on their current working directory.
-export PYTHONPATH="${PYTHONPATH:-..}"
+# Export Python path
+export PYTHONPATH=".."

 ###############################################################################
 # Helper Functions
@@ -54,108 +52,6 @@ cleanup_network() {
  fi
 }

-prepare_artifact_image() {
-  if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then
-    return 1
-  fi
-  if ! command -v buildkite-agent >/dev/null 2>&1; then
-    echo "buildkite-agent not found; cannot download ROCm wheel artifact"
-    return 1
-  fi
-
-  local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}"
-  local archive=""
-  local metadata_file=""
-  local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}"
-  local artifact_image=""
-  local artifact_key=""
-  local base_digest=""
-  local wheel_dir=""
-  local context_dir=""
-  local workspace_dir=""
-
-  artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX)
-  wheel_dir="${artifact_work_dir}/wheels"
-  context_dir="${artifact_work_dir}/context"
-  workspace_dir="${context_dir}/workspace"
-  mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}"
-
-  echo "--- Downloading ROCm wheel artifact"
-  if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then
-    echo "Failed to download ${artifact_glob}"
-    return 1
-  fi
-  buildkite-agent artifact download \
-    "artifacts/vllm-rocm-install/ci-base-image.txt" \
-    "${artifact_work_dir}" >/dev/null 2>&1 || true
-
-  archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1)
-  if [[ -z "${archive}" || ! -f "${archive}" ]]; then
-    echo "ROCm wheel artifact archive was not found"
-    return 1
-  fi
-
-  metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1)
-  if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then
-    base_image=$(tr -d '[:space:]' < "${metadata_file}")
-  fi
-
-  echo "--- Preparing local ROCm test image"
-  echo "Base image: ${base_image}"
-  docker pull "${base_image}" || return 1
-  base_digest=$(
-    docker image inspect \
-      --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \
-      "${base_image}" 2>/dev/null || printf '%s' "${base_image}"
-  )
-
-  artifact_key=$(
-    {
-      printf 'base-image:%s\n' "${base_digest}"
-      sha256sum "${archive}"
-    } | sha256sum | cut -c1-24
-  )
-  artifact_image="rocm/vllm-ci-artifact:${artifact_key}"
-
-  if docker image inspect "${artifact_image}" >/dev/null 2>&1; then
-    echo "Using existing local ROCm artifact image: ${artifact_image}"
-    image_name="${artifact_image}"
-    return 0
-  fi
-
-  tar -xzf "${archive}" -C "${wheel_dir}" || return 1
-  if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then
-    echo "ROCm wheel artifact did not contain a wheel"
-    return 1
-  fi
-  if [[ ! -d "${wheel_dir}/tests" ]]; then
-    echo "ROCm wheel artifact did not contain the test workspace"
-    return 1
-  fi
-
-  cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1
-  tar -C "${wheel_dir}" --exclude='*.whl' -cf - . \
-    | tar -C "${workspace_dir}" -xf - || return 1
-  cat > "${context_dir}/Dockerfile" <<'EOF'
-ARG BASE_IMAGE
-FROM ${BASE_IMAGE}
-COPY wheels/ /tmp/vllm-wheels/
-COPY workspace/ /vllm-workspace/
-RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \
-    && rm -rf /tmp/vllm-wheels
-WORKDIR /vllm-workspace
-EOF
-
-  echo "--- Building local ROCm test image"
-  docker build \
-    --pull=false \
-    --build-arg "BASE_IMAGE=${base_image}" \
-    -t "${artifact_image}" \
-    "${context_dir}" || return 1
-  image_name="${artifact_image}"
-  return 0
-}
-
 is_multi_node() {
  local cmds="$1"
  # Primary signal: NUM_NODES environment variable set by the pipeline
@@ -347,30 +243,22 @@ report_docker_usage

 # --- Pull test image ---
 echo "--- Pulling container"
-image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}"
-artifact_work_dir=""
+image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+docker pull "${image_name}"

 remove_docker_container() {
-  if docker container inspect "${container_name}" >/dev/null 2>&1; then
-    docker rm -f "${container_name}" || true
-  fi
-  if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then
-    docker image rm -f "${image_name}" || true
-  else
-    # Keep images by default so later jobs on the same AMD node can reuse layers.
-    echo "Keeping ROCm test image locally: ${image_name}"
-  fi
-  if [[ -n "${artifact_work_dir}" ]]; then
-    rm -rf "${artifact_work_dir}"
-  fi
+  # docker run uses --rm, so the container is normally already gone when the
+  # EXIT trap runs. Cleanup is best-effort and must not affect the test result.
+  docker rm -f "${container_name}" >/dev/null 2>&1 || true
 }
-trap remove_docker_container EXIT

-if ! prepare_artifact_image; then
-  echo "Using full ROCm CI image: ${image_name}"
-  docker pull "${image_name}" || exit 1
-fi
+on_exit() {
+  local exit_code=$?
+  remove_docker_container
+  exit "$exit_code"
+}
+trap on_exit EXIT

 # --- Prepare commands ---
 echo "--- Running container"
@@ -379,14 +267,6 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"

-# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm
-# CI image currently raises downloads to 60s. AMD model-test jobs routinely
-# start from a cold or partially-populated shared cache, and the 60s read cap
-# has still timed out before pytest reached the vLLM behavior under test.
-# Keep the CI default explicit and overridable from the Buildkite environment.
-: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-: "${HF_HUB_ETAG_TIMEOUT:=60}"
-
 # ---- Command source selection ----
 # Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
 # Fall back to $* for backward compatibility, but warn that inner
@@ -426,14 +306,7 @@ fi

 echo "Final commands: $commands"

-MYPYTHONPATH="/vllm-workspace"
-
-container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
-container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}"
-container_job_id_short="${container_job_id:0:8}"
-CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}"
-CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache"
-CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\""
+MYPYTHONPATH=".."

 # Verify GPU access
 render_gid=$(getent group render | cut -d: -f3)
@@ -510,8 +383,6 @@ else
    --group-add "$render_gid" \
    --rm \
    -e HF_TOKEN \
-    -e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
-    -e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
    -e AWS_ACCESS_KEY_ID \
    -e AWS_SECRET_ACCESS_KEY \
    -e BUILDKITE_PARALLEL_JOB \
@@ -519,15 +390,10 @@ else
    -v "${HF_CACHE}:${HF_MOUNT}" \
    -e "HF_HOME=${HF_MOUNT}" \
    -e "PYTHONPATH=${MYPYTHONPATH}" \
-    -e "TMPDIR=${CONTAINER_TMPDIR}/tmp" \
-    -e "TORCHINDUCTOR_CACHE_DIR=${CONTAINER_CACHE_ROOT}/torchinductor" \
-    -e "TRITON_CACHE_DIR=${CONTAINER_CACHE_ROOT}/triton" \
-    -e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
-    -e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
    -e "PYTORCH_ROCM_ARCH=" \
    --name "${container_name}" \
    "${image_name}" \
-    /bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
+    /bin/bash -c "${commands}"

  exit_code=$?
  handle_pytest_exit "$exit_code"
@@ -324,6 +324,23 @@ IMAGE="${IMAGE_TAG_XPU:-${image_name}}"

 echo "Using image: ${IMAGE}"

+if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
+  echo "Image already exists locally, skipping pull"
+else
+  echo "Image not found locally, waiting for lock..."
+
+  flock /tmp/docker-pull.lock bash -c "
+    if docker image inspect '${IMAGE}' >/dev/null 2>&1; then
+      echo 'Image already pulled by another runner'
+    else
+      echo 'Pulling image...'
+      timeout 900 docker pull '${IMAGE}'
+    fi
+  "
+
+  echo "Pull step completed"
+fi
+
 remove_docker_container() {
  docker rm -f "${container_name}" || true
 }
@@ -340,12 +357,9 @@ export HF_TOKEN ZE_AFFINITY_MASK

 {
  flock 9
-  if docker image inspect "${IMAGE}" >/dev/null 2>&1; then
-    echo "Image already exists locally, skipping pull"
-  else
-    echo "Image not found locally, pulling image..."
+  if ! docker image inspect "${IMAGE}" >/dev/null 2>&1; then
+    echo 'Image missing before container creation, pulling again...'
    timeout 900 docker pull "${IMAGE}"
-    echo "Pull step completed"
  fi

  docker create \
@@ -358,8 +372,6 @@ export HF_TOKEN ZE_AFFINITY_MASK
    --entrypoint='' \
    -e HF_TOKEN \
    -e ZE_AFFINITY_MASK \
-    -e BUILDKITE_PARALLEL_JOB \
-    -e BUILDKITE_PARALLEL_JOB_COUNT \
    -e CMDS \
    --name "${container_name}" \
    "${IMAGE}" \
@@ -110,36 +110,6 @@ install_uv() {
    | env UV_INSTALL_DIR="$CARGO_HOME/bin" sh
 }

-setup_pyo3_python() {
-  local python_version="${PYO3_PYTHON_VERSION:-3.12}"
-
-  log_section "Installing Python ${python_version} for PyO3 tests"
-  uv python install "$python_version"
-  PYO3_PYTHON="$(uv python find \
-    --managed-python \
-    --no-project \
-    --resolve-links \
-    "$python_version")"
-  export PYO3_PYTHON
-
-  local python_libdir
-  python_libdir="$("$PYO3_PYTHON" - <<'PY'
-import pathlib
-import sysconfig
-
-libdir = pathlib.Path(sysconfig.get_config_var("LIBDIR"))
-ldlibrary = sysconfig.get_config_var("LDLIBRARY")
-assert sysconfig.get_config_var("Py_ENABLE_SHARED") == 1
-assert ldlibrary
-assert (libdir / ldlibrary).exists(), libdir / ldlibrary
-print(libdir)
-PY
-)"
-
-  export LD_LIBRARY_PATH="${python_libdir}:${LD_LIBRARY_PATH:-}"
-  export LIBRARY_PATH="${python_libdir}:${LIBRARY_PATH:-}"
-}
-
 run_style_clippy() {
  install_cargo_sort

@@ -162,7 +132,6 @@ run_style_clippy() {

 run_tests() {
  install_uv
-  setup_pyo3_python
  install_cargo_nextest

  log_section "Running cargo nextest"
@@ -88,16 +88,16 @@
 #     - Do NOT remove `VLLM_WORKER_MULTIPROC_METHOD=spawn` setting as ROCm requires this for certain models to function.            #
 #   * [Transformers Nightly Models]: Whisper needs `VLLM_WORKER_MULTIPROC_METHOD=spawn` to avoid deadlock.                          #
 #   * [Plugin Tests (2 GPUs)]:                                                                                                      #
-#     - {`pytest -v -s plugins_tests/test_oot_registration_online.py`}: It needs a clean process                                    #
-#     - {`pytest -v -s plugins_tests/test_oot_registration_offline.py`}: It needs a clean process                                   #
-#     - {`pytest -v -s plugins_tests/lora_resolvers`}: Unit tests for in-tree lora resolver plugins                                 #
+#     - {`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                                       #
 #   * [LoRA TP (Distributed)]:                                                                                                      #
 #     - There is some Tensor Parallelism related processing logic in LoRA that requires multi-GPU testing for validation.           #
 #     - {`pytest -v -s -x lora/test_gptoss_tp.py`}: Disabled for now because MXFP4 backend on non-cuda platform doesn't support     #
 #                                                   LoRA yet.                                                                       #
 #   * [Distributed Tests (NxGPUs)(HW-TAG)]: Don't test llama model here, it seems hf implementation is buggy. See:                  #
 #                                           https://github.com/vllm-project/vllm/pull/5689                                          #
-#   * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293      #
+#   * [Distributed Tests (NxGPUs)(HW-TAG)]: Some old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293     #
 #                                           in favor of new tests in fusions_e2e. We avoid replicating the new jobs in              #
 #                                           this file as it's deprecated.                                                           #
 #                                                                                                                                   #
@@ -315,6 +315,24 @@ steps:
  - pytest -v -s distributed/test_pp_cudagraph.py
  - pytest -v -s distributed/test_pipeline_parallel.py

+#----------------------------------------------------------  mi250 · engine  -----------------------------------------------------------#
+
+- label: Engine # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/engine
+  - tests/test_sequence
+  - tests/test_config
+  - tests/test_logger
+  - tests/test_vllm_port
+  commands:
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+
 #-----------------------------------------------------------  mi250 · evals  -----------------------------------------------------------#

 - label: Multi-Modal Accuracy Eval (Small Models) # TBD
@@ -431,6 +449,29 @@ steps:
  commands:
  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py

+#------------------------------------------------------  mi250 · model_executor  -------------------------------------------------------#
+
+- label: Model Executor # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
+  agent_pool: mi250_1
+  optional: true
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/engine/arg_utils.py
+  - vllm/config/model.py
+  - vllm/model_executor
+  - tests/model_executor
+  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - apt-get update && apt-get install -y curl libsodium23
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s model_executor -m '(not slow_test)'
+  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
+
 #------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#

 - label: Basic Models Test (Other CPU) # TBD
@@ -617,9 +658,9 @@ 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 plugins_tests/test_oot_registration_online.py # it needs a clean process
-  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
-  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+  - pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s models/test_oot_registration.py
+  - pytest -v -s plugins/lora_resolvers

 #------------------------------------------------------------  mi250 · v1  -------------------------------------------------------------#

@@ -804,7 +845,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: V1 e2e (2 GPUs) # TBD
  timeout_in_minutes: 180
@@ -830,7 +871,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#

@@ -890,10 +931,10 @@ steps:
  - vllm/
  - tests/basic_correctness/test_basic_correctness
  - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_mem.py
+  - tests/basic_correctness/test_cumem.py
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_mem.py
+  - pytest -v -s basic_correctness/test_cumem.py
  - pytest -v -s basic_correctness/test_basic_correctness.py
  - pytest -v -s basic_correctness/test_cpu_offload.py

@@ -1187,39 +1228,7 @@ steps:

 #--------------------------------------------------------  mi300 · entrypoints  --------------------------------------------------------#

- label: Entrypoints Unit Tests # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  fast_check: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/entrypoints
-  - tests/entrypoints/unit_tests
-  - tests/entrypoints/weight_transfer
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s entrypoints/unit_tests
-  - pytest -v -s entrypoints/weight_transfer
-
- label: Entrypoints Integration (LLM) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/llm
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
-  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
-
- label: Entrypoints Integration (API Server) # TBD
+- label: Entrypoints Integration (API Server 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
@@ -1235,7 +1244,7 @@ steps:
  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc

- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
+- label: Entrypoints Integration (API Server openai - Part 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
@@ -1249,59 +1258,44 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/ --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py

- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
+- label: Entrypoints Integration (API Server openai - Part 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  torch_nightly: true
-  optional: true
  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/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-
- label: Entrypoints Integration (API Server Generate) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/tool_use
-  - tests/entrypoints/tool_parsers
-  - tests/entrypoints/anthropic
  - tests/entrypoints/generate
+  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s tool_use
-  - pytest -v -s entrypoints/tool_parsers
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/test_chat_utils.py
  - pytest -v -s entrypoints/generate
-  - pytest -v -s entrypoints/anthropic
+  - pytest -v -s tool_use

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

 - label: Entrypoints Integration (Speech to Text) # TBD
  timeout_in_minutes: 180
@@ -1317,19 +1311,23 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/speech_to_text

- label: Entrypoints Integration (Multimodal)
+- label: Entrypoints Integration (LLM) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
-  - tests/entrypoints/multimodal
+  - tests/entrypoints/llm
+  - tests/entrypoints/offline_mode
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/multimodal
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s entrypoints/llm/test_generate.py
+  - pytest -v -s entrypoints/offline_mode

 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
@@ -1345,15 +1343,43 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/pooling

+- label: Entrypoints Integration (Responses API) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  fast_check: true
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai/responses
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/openai/responses
+
+- label: Entrypoints Unit Tests # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  fast_check: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/entrypoints
+  - tests/entrypoints/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s entrypoints/openai/tool_parsers
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
+
 - label: OpenAI API correctness # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/
  - vllm/entrypoints/openai/
+  - vllm/model_executor/models/whisper.py
  - vllm/model_executor/layers/
  - vllm/v1/attention/backends/
  - vllm/v1/attention/selector.py
@@ -1407,7 +1433,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
  - csrc/
@@ -1714,29 +1739,6 @@ steps:
  - pytest -v -s -x lora/test_gptoss_tp.py
  - pytest -v -s -x lora/test_qwen35_densemodel_lora.py

-#------------------------------------------------------  mi300 · model_executor  -------------------------------------------------------#
-
- label: Model Executor # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/engine/arg_utils.py
-  - vllm/config/model.py
-  - vllm/model_executor
-  - tests/model_executor
-  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - apt-get update && apt-get install -y curl libsodium23
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s model_executor -m '(not slow_test)'
-  - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
-
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#

 - label: Language Models Test (Extended Pooling)  # TBD
@@ -2182,72 +2184,10 @@ steps:
  commands:
  - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"

- label: Speculators Correctness # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/config/speculative.py
-  - vllm/engine/arg_utils.py
-  - vllm/transformers_utils/config.py
-  - vllm/transformers_utils/configs/speculators/
-  - vllm/v1/spec_decode/
-  - vllm/v1/worker/gpu/spec_decode/
-  - vllm/v1/worker/gpu_model_runner.py
-  - vllm/v1/sample/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/layers/
-  - vllm/model_executor/models/llama_eagle3.py
-  - vllm/model_executor/models/qwen3.py
-  - vllm/model_executor/models/qwen3_dflash.py
-  - vllm/model_executor/models/registry.py
-  - vllm/_aiter_ops.py
-  - tests/evals/gsm8k/
-  - tests/v1/spec_decode/test_speculators_correctness.py
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-  - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
-
- label: Extract Hidden States Integration # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/config/speculative.py
-  - vllm/distributed/kv_transfer/kv_connector/
-  - vllm/model_executor/layers/attention/
-  - vllm/model_executor/layers/mamba/
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/models/extract_hidden_states.py
-  - vllm/model_executor/models/llama.py
-  - vllm/model_executor/models/qwen3_5.py
-  - vllm/model_executor/models/qwen3_next.py
-  - vllm/model_executor/models/registry.py
-  - vllm/transformers_utils/configs/extract_hidden_states.py
-  - vllm/transformers_utils/configs/qwen3_5.py
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - vllm/v1/kv_cache_interface.py
-  - vllm/v1/spec_decode/extract_hidden_states.py
-  - vllm/v1/worker/gpu_model_runner.py
-  - vllm/_aiter_ops.py
-  - tests/v1/kv_connector/extract_hidden_states_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s v1/kv_connector/extract_hidden_states_integration
-
 - label: V1 attention (H100-MI300) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/config/attention.py
@@ -2406,7 +2346,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - CROSS_LAYERS_BLOCKS=True ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: Distributed DP Tests (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -2442,7 +2382,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -2456,7 +2396,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -2470,7 +2410,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - HYBRID_SSM=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: V1 e2e (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -2606,7 +2546,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2616,7 +2555,7 @@ steps:
  - tests/test_logger
  - tests/test_vllm_port
  commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py

 #-----------------------------------------------------------  mi325 · evals  -----------------------------------------------------------#

@@ -2698,7 +2637,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2798,7 +2736,7 @@ steps:

 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#

- label: Entrypoints Integration (API Server) # TBD
+- label: Entrypoints Integration (API Server 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
@@ -2814,7 +2752,7 @@ steps:
  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc

- label: Entrypoints Integration (API Server OpenAI - Part 1) # TBD
+- label: Entrypoints Integration (API Server openai - Part 1) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
@@ -2828,9 +2766,9 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py

- label: Entrypoints Integration (API Server OpenAI - Part 2) # TBD
+- label: Entrypoints Integration (API Server openai - Part 2) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
@@ -2842,31 +2780,29 @@ steps:
  - vllm/
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-
- label: Entrypoints Integration (API Server Generate) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-  agent_pool: mi355_1
-  fast_check: true
-  torch_nightly: true
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/tool_use
-  - tests/entrypoints/tool_parsers
-  - tests/entrypoints/anthropic
  - tests/entrypoints/generate
+  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s tool_use
-  - pytest -v -s entrypoints/tool_parsers
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/test_chat_utils.py
  - pytest -v -s entrypoints/generate
-  - pytest -v -s entrypoints/anthropic
+  - pytest -v -s tool_use
+
+- label: Entrypoints Integration (API Server openai - Part 3) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  agent_pool: mi355_1
+  fast_check: true
+  torch_nightly: true
+  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/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (Speech to Text) # TBD
  timeout_in_minutes: 180
@@ -2882,20 +2818,6 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/speech_to_text

- label: Entrypoints Integration (Multimodal)
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
-  agent_pool: mi355_1
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/multimodal
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/multimodal
-
 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -3073,7 +2995,7 @@ steps:
  - vllm/_aiter_ops.py
  commands:
  - rocm-smi
-  - python3 examples/basic/offline_inference/chat.py --attention-backend TRITON_ATTN
+  - python3 examples/basic/offline_inference/chat.py
  - pytest -v -s tests/kernels/attention/test_attention_selector.py

 - label: Kernels Attention Test %N # TBD
@@ -3437,7 +3359,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -3452,7 +3374,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -3467,7 +3389,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+  - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

 #------------------------------------------------------  mi355 · weight_loading  -------------------------------------------------------#

@@ -2,8 +2,8 @@ group: Attention
 depends_on: 
  - image-build
 steps:
- label: V1 attention (H100-MI300)
-  key: v1-attention-h100-mi300
+- label: V1 attention (H100)
+  key: v1-attention-h100
  timeout_in_minutes: 30
  device: h100
  source_file_dependencies:
@@ -13,20 +13,6 @@ steps:
    - tests/v1/attention
  commands:
    - pytest -v -s v1/attention
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 70
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/config/attention.py
-      - vllm/model_executor/layers/attention
-      - vllm/v1/attention
-      - tests/v1/attention
-      - vllm/_aiter_ops.py
-      - vllm/envs.py
-      - vllm/platforms/rocm.py

 - label: V1 attention (B200)
  key: v1-attention-b200
@@ -10,9 +10,9 @@ steps:
  - vllm/
  - tests/basic_correctness/test_basic_correctness
  - tests/basic_correctness/test_cpu_offload
-  - tests/basic_correctness/test_mem.py
+  - tests/basic_correctness/test_cumem.py
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_mem.py
+  - pytest -v -s basic_correctness/test_cumem.py
  - pytest -v -s basic_correctness/test_basic_correctness.py
  - pytest -v -s basic_correctness/test_cpu_offload.py
@@ -26,12 +26,6 @@ steps:
  - tests/test_jit_monitor.py
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd

 - label: Engine (1 GPU)
  key: engine-1-gpu
@@ -44,7 +38,7 @@ steps:
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      timeout_in_minutes: 40
      depends_on:
      - image-build-amd
@@ -8,11 +8,10 @@ steps:
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/entrypoints
-  - tests/entrypoints/unit_tests
-  - tests/entrypoints/weight_transfer
+  - tests/entrypoints/
  commands:
-  - pytest -v -s entrypoints/unit_tests
-  - pytest -v -s entrypoints/weight_transfer
+  - pytest -v -s entrypoints/openai/tool_parsers
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate

 - label: Entrypoints Integration (LLM)
  key: entrypoints-integration-llm
@@ -21,20 +20,84 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/llm
+  - tests/entrypoints/offline_mode
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
+  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
      device: mi325_1
+      soft_fail: true
      depends_on:
      - image-build-amd

- label: Entrypoints Integration (API Server)
-  key: entrypoints-integration-api-server
+- label: Entrypoints Integration (API Server openai - Part 1)
+  key: entrypoints-integration-api-server-openai-part-1
+  timeout_in_minutes: 50
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai
+  - 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_oot_registration.py
+  mirror:
+    amd:
+      device: mi325_1
+      soft_fail: true
+      timeout_in_minutes: 80
+      depends_on:
+      - image-build-amd
+
+- label: Entrypoints Integration (API Server openai - Part 2)
+  key: entrypoints-integration-api-server-openai-part-2
+  timeout_in_minutes: 50
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai
+  - tests/entrypoints/test_chat_utils
+  - tests/entrypoints/generate
+  - tests/tool_use
+  commands:
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/test_chat_utils.py
+  - pytest -v -s entrypoints/generate
+  - pytest -v -s tool_use
+  mirror:
+    amd:
+      device: mi325_1
+      soft_fail: true
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
+
+- label: Entrypoints Integration (API Server openai - Part 3)
+  key: entrypoints-integration-api-server-openai-part-3
+  timeout_in_minutes: 50
+  device: h200_18gb
+  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/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
+  mirror:
+    amd:
+      device: mi325_1
+      soft_fail: true
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
+
+- label: Entrypoints Integration (API Server 2)
  device: h200_35gb
+  key: entrypoints-integration-api-server-2
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -46,79 +109,11 @@ steps:
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
      depends_on:
      - image-build-amd

- label: Entrypoints Integration (API Server OpenAI - Part 1)
-  key: entrypoints-integration-api-server-openai-part-1
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  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/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 80
-      depends_on:
-      - image-build-amd
-
- label: Entrypoints Integration (API Server OpenAI - Part 2)
-  key: entrypoints-integration-api-server-openai-part-2
-  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/chat_completion
-  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 80
-      depends_on:
-      - image-build-amd
-
- label: Entrypoints Integration (API Server Generate)
-  key: entrypoints-integration-api-server-generate
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/tool_use
-  - tests/entrypoints/tool_parsers
-  - tests/entrypoints/anthropic
-  - tests/entrypoints/generate
-  commands:
-  - pytest -v -s tool_use
-  - pytest -v -s entrypoints/tool_parsers
-  - pytest -v -s entrypoints/generate
-  - pytest -v -s entrypoints/anthropic
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
-
- label: Entrypoints Integration (Responses API)
-  key: entrypoints-integration-responses-api
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/openai/responses
-  commands:
-  - pytest -v -s entrypoints/openai/responses
-
 - label: Entrypoints Integration (Speech to Text)
  device: h200_35gb
  key: entrypoints-integration-speech_to_text
@@ -131,18 +126,6 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/speech_to_text

- label: Entrypoints Integration (Multimodal)
-  device: h200_35gb
-  key: entrypoints-integration-multimodal
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/multimodal
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/multimodal
-
 - label: Entrypoints Integration (Pooling)
  key: entrypoints-integration-pooling
  timeout_in_minutes: 50
@@ -154,6 +137,16 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/pooling

+- label: Entrypoints Integration (Responses API)
+  key: entrypoints-integration-responses-api
+  timeout_in_minutes: 50
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/openai/responses
+  commands:
+  - pytest -v -s entrypoints/openai/responses
+
 - label: OpenAI API Correctness
  key: openai-api-correctness
  timeout_in_minutes: 30
@@ -163,20 +156,3 @@ steps:
  - vllm/entrypoints/openai/
  commands: # LMEval
  - pytest -s entrypoints/openai/correctness/
-  mirror:
-    amd:
-      device: mi300_1
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - csrc/
-      - vllm/entrypoints/openai/
-      - vllm/model_executor/layers/
-      - vllm/v1/attention/backends/
-      - vllm/v1/attention/selector.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py
-      - vllm/model_executor/model_loader/
-      commands:
-      - bash ../tools/install_torchcodec_rocm.sh || exit 1
-      - pytest -s entrypoints/openai/correctness/
@@ -21,9 +21,8 @@ steps:
  - csrc/
  - tests/kernels/core
  - tests/kernels/test_concat_mla_q.py
-  - tests/kernels/test_fused_qk_norm_rope_gate.py
  commands:
-    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py
+    - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py

 - label: Kernels MiniMax Reduce RMS Test (2 GPUs)
  key: kernels-minimax-reduce-rms-test-2-gpus
@@ -87,7 +86,7 @@ steps:
  parallelism: 2
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
@@ -300,4 +299,3 @@ steps:
  - vllm/config
  commands:
    - pytest -v -s kernels/moe/test_moe_layer.py
-    - pytest -v -s kernels/moe/test_deepep_v2_moe.py
@@ -12,21 +12,6 @@ steps:
  autorun_on_main: true
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 55
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - csrc/
-      - vllm/model_executor/layers/quantization
-      - vllm/model_executor/models/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/attention/backends/
-      - vllm/v1/attention/selector.py
-      - vllm/_aiter_ops.py
-      - vllm/platforms/rocm.py

 # - label: LM Eval Large Models (4 GPUs)(A100)
 #   device: a100
@@ -281,7 +281,6 @@ steps:
  - vllm/model_executor/layers/quantization/quark/
  - vllm/multimodal/
  - vllm/outputs.py
-  - vllm/parser/
  - vllm/platforms/
  - vllm/pooling_params.py
  - vllm/ray/
@@ -300,9 +299,9 @@ steps:
  - tests/multimodal
  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
+  - tests/tokenizers_
  - tests/reasoning
  - tests/tool_parsers
-  - tests/tokenizers_
  - tests/parser
  - tests/transformers_utils
  - tests/config
@@ -315,9 +314,9 @@ steps:
  - pytest -v -s test_ray_env.py
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
+  - pytest -v -s tokenizers_
  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
-  - pytest -v -s tokenizers_
  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config
@@ -94,13 +94,11 @@ steps:
    - vllm/v1/worker/gpu_worker.py
    - tests/distributed/test_pipeline_parallel.py
    - tests/distributed/test_pp_cudagraph.py
-    - tests/v1/distributed/test_pp_dp_v2.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"
-    - pytest -v -s v1/distributed/test_pp_dp_v2.py

 - label: Model Runner V2 Spec Decode
  device: h200_35gb
@@ -51,7 +51,6 @@ steps:
    torch_nightly: {}
    amd:
      device: mi325_1
-      timeout_in_minutes: 90
      depends_on:
      - image-build-amd
      commands:
@@ -15,7 +15,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      depends_on:
      - image-build-amd

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

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

@@ -153,12 +153,3 @@ steps:
  - tests/models/multimodal/pooling
  commands:
    - pytest -v -s models/multimodal/pooling -m 'not core_model'
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/
-      - tests/models/multimodal/pooling
@@ -37,6 +37,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 plugins_tests/test_oot_registration_online.py # it needs a clean process
-  - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
-  - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+  - pytest -v -s entrypoints/openai/chat_completion/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
@@ -37,21 +37,6 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - vllm/transformers_utils/configs/speculators/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py

 - label: Spec Decode Speculators + MTP Nightly B200
  key: spec-decode-speculators-mtp-nightly-b200
@@ -76,20 +61,6 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 65
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py

 - label: Spec Decode Draft Model
  key: spec-decode-draft-model
@@ -101,20 +72,6 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 50
-      depends_on:
-      - image-build-amd
-      source_file_dependencies:
-      - vllm/v1/spec_decode/
-      - vllm/v1/worker/gpu/spec_decode/
-      - vllm/model_executor/model_loader/
-      - vllm/v1/sample/
-      - vllm/model_executor/layers/
-      - tests/v1/e2e/spec_decode/
-      - vllm/platforms/rocm.py

 - label: Spec Decode Draft Model Nightly B200
  key: spec-decode-draft-model-nightly-b200
@@ -23,14 +23,9 @@

 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
-/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @yewentao256 @ProExpertProg
+/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
 /vllm/config/cache.py @heheda12345

-# Config utils
-/vllm/config/utils.py @hmellor
-/vllm/engine/arg_utils.py @hmellor
-/vllm/utils/argparse_utils.py
-
 # Entrypoints
 /vllm/entrypoints/anthropic @mgoin @DarkLight1337
 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
@@ -39,11 +34,10 @@
 /vllm/entrypoints/speech_to_text/realtime @njhill
 /vllm/entrypoints/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
-/vllm/entrypoints/serve/sagemaker @DarkLight1337
+/vllm/entrypoints/sagemaker @DarkLight1337
 /vllm/entrypoints/serve @njhill
 /vllm/entrypoints/*.py @njhill
 /vllm/entrypoints/chat_utils.py @DarkLight1337
-/vllm/entrypoints/offline_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337

 # Rust Frontend
@@ -21,9 +21,6 @@ pull_request_rules:
    - check-failure=pre-commit
    - -closed
    - -draft
-    - or:
-      - label=ready
-      - label=verified
  actions:
    comment:
      message: |
@@ -39,6 +36,18 @@ pull_request_rules:

        For future commits, `pre-commit` will run automatically on changed files before each commit.

+        > [!TIP]
+        > <details>
+        > <summary>Is <code>mypy</code> failing?</summary>
+        > <br/>
+        > <code>mypy</code> is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
+        >
+        > ```bash
+        > # For mypy (substitute "3.10" with the failing version if needed)
+        > pre-commit run --hook-stage manual mypy-3.10
+        > ```
+        > </details>
+
 - name: comment-dco-failure
  description: Comment on PR when DCO check fails
  conditions:
@@ -49,7 +49,7 @@ jobs:
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
-    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
+    - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
      with:
        python-version: "3.12"
    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
@@ -15,7 +15,7 @@ jobs:
      actions: write
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
+      - uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
        with:
          # Increasing this value ensures that changes to this workflow
          # propagate to all issues and PRs in days rather than months
@@ -148,27 +148,33 @@ repos:
    language: python
    entry: python tools/pre_commit/generate_nightly_torch_test.py
    files: ^requirements/test/cuda\.(in|txt)$
-  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
-    name: Run mypy for Python 3.10
-    entry: python tools/pre_commit/mypy.py "3.10"
+  - id: mypy-local
+    name: Run mypy locally for lowest supported Python version
+    entry: python tools/pre_commit/mypy.py 0 "3.10"
+    stages: [pre-commit] # Don't run in CI
    <<: &mypy_common
      language: python
      types_or: [python, pyi]
      require_serial: true
-      additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+      additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
+    name: Run mypy for Python 3.10
+    entry: python tools/pre_commit/mypy.py 1 "3.10"
+    <<: *mypy_common
+    stages: [manual] # Only run in CI
  - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
    name: Run mypy for Python 3.11
-    entry: python tools/pre_commit/mypy.py "3.11"
+    entry: python tools/pre_commit/mypy.py 1 "3.11"
    <<: *mypy_common
    stages: [manual] # Only run in CI
  - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
    name: Run mypy for Python 3.12
-    entry: python tools/pre_commit/mypy.py "3.12"
+    entry: python tools/pre_commit/mypy.py 1 "3.12"
    <<: *mypy_common
    stages: [manual] # Only run in CI
  - id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
    name: Run mypy for Python 3.13
-    entry: python tools/pre_commit/mypy.py "3.13"
+    entry: python tools/pre_commit/mypy.py 1 "3.13"
    <<: *mypy_common
    stages: [manual] # Only run in CI
  - id: shellcheck
@@ -98,13 +98,11 @@ pre-commit run --all-files
 pre-commit run ruff-check --all-files

 # Run mypy as it is in CI:
-pre-commit run mypy-3.12 --all-files --hook-stage manual
+pre-commit run mypy-3.10 --all-files --hook-stage manual
 ```

 The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.

-Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
-
 ### Commit messages

 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -112,8 +112,6 @@ endif()
 #
 # spinloop extension (pure CXX; must stay above the non-CUDA device branch so
 # CPU builds define the target before the early return)
-# This extension requires SABI 3.11 since it relies on Py_buffer support. Loading
-# failure is handled gracefully on vLLM side for lower Python versions.
 #
 set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
 set(SPINLOOP_COMPILE_FLAGS "")
@@ -307,10 +305,18 @@ endif()
 #

 set(VLLM_EXT_SRC
+  "csrc/cuda_view.cu"
+  "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
  "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
+  "csrc/cuda_utils_kernels.cu"
+  "csrc/custom_all_reduce.cu"
+  "csrc/torch_bindings.cpp"
+  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
+  list(APPEND VLLM_EXT_SRC
+    "csrc/minimax_reduce_rms_kernel.cu")
+
  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
@@ -343,6 +349,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  endif()
  FetchContent_MakeAvailable(cutlass)

+  list(APPEND VLLM_EXT_SRC
+    "csrc/cutlass_extensions/common.cpp")
+
  set_gencode_flags_for_srcs(
    SRCS "${VLLM_EXT_SRC}"
    CUDA_ARCHS "${CUDA_ARCHS}")
@@ -494,12 +503,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
    set(SRCS
-      "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
+      "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
+      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
      CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_EXT_SRC "${SRCS}")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
    message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
  else()
@@ -589,12 +598,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 endif()

 if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
+  # Add QuickReduce kernels
  list(APPEND VLLM_EXT_SRC
    "csrc/custom_quickreduce.cu"
-    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
  )
 # if ROCM endif
 endif()
@@ -627,9 +633,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
-    "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
-    "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
-    "csrc/libtorch_stable/permute_cols.cu"
    "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
    "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
    "csrc/libtorch_stable/pos_encoding_kernels.cu"
@@ -637,7 +640,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/layernorm_kernels.cu"
    "csrc/libtorch_stable/layernorm_quant_kernels.cu"
    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
-    "csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
    "csrc/libtorch_stable/attention/merge_attn_states.cu"
    "csrc/libtorch_stable/sampler.cu"
    "csrc/libtorch_stable/topk.cu"
@@ -645,21 +647,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
    "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels_fused.cu"
-    "csrc/libtorch_stable/custom_all_reduce.cu"
-    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/libtorch_stable/cache_kernels_fused.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA")
    list(APPEND VLLM_STABLE_EXT_SRC
-      "csrc/libtorch_stable/cuda_view.cu"
-      "csrc/libtorch_stable/cuda_utils_kernels.cu"
-      "csrc/libtorch_stable/cutlass_extensions/common.cpp"
+      "csrc/cuda_utils_kernels.cu"
+      "csrc/cutlass_extensions/common.cpp"
      "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
-      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
-      "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
+      "csrc/libtorch_stable/permute_cols.cu"
+      "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
+      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
+      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")

    set_gencode_flags_for_srcs(
      SRCS "${VLLM_STABLE_EXT_SRC}"
@@ -1072,25 +1072,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    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/hip APIs from C-shim
  if(VLLM_GPU_LANG STREQUAL "CUDA")
-    # Set TORCH_TARGET_VERSION for stable ABI compatibility.
-    # This ensures we only use C-shim APIs available in PyTorch 2.11.
-    # _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
-    # which is currently set to 2.11.
-    target_compile_definitions(_C_stable_libtorch PRIVATE
-      TORCH_TARGET_VERSION=0x020B000000000000ULL)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
  elseif(VLLM_GPU_LANG STREQUAL "HIP")
-    # 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)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
  endif()

@@ -1308,8 +1303,7 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    set(VLLM_ROCM_HAS_GFX1100 ON)
    list(APPEND VLLM_ROCM_EXT_SRC
      "csrc/rocm/q_gemm_rdna3.cu"
-      "csrc/rocm/q_gemm_rdna3_wmma.cu"
-      "csrc/rocm/moe_q_gemm_rdna3.cu")
+      "csrc/rocm/q_gemm_rdna3_wmma.cu")
  endif()

  define_extension_target(
@@ -4,7 +4,6 @@ include requirements/cuda.txt
 include requirements/rocm.txt
 include requirements/cpu.txt
 include CMakeLists.txt
-include tools/build_rust.py

 recursive-include cmake *
 recursive-include csrc *
@@ -0,0 +1,143 @@
+#!/bin/bash
+
+# benchmark the overhead of disaggregated prefill.
+# methodology:
+# - send all request to prefill vLLM instance. It will buffer KV cache.
+# - then send all request to decode instance.
+# - The TTFT of decode instance is the overhead.
+
+set -ex
+
+kill_gpu_processes() {
+  # kill all processes on GPU.
+  pgrep pt_main_thread | xargs -r kill -9
+  pgrep python3 | xargs -r kill -9
+  # vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
+  pgrep VLLM | xargs -r kill -9
+  sleep 10
+
+  # remove vllm config file
+  rm -rf ~/.config/vllm
+
+  # Print the GPU memory usage
+  # so that we know if all GPU processes are killed.
+  gpu_memory_usage=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0)
+  # The memory usage should be 0 MB.
+  echo "GPU 0 Memory Usage: $gpu_memory_usage MB"
+}
+
+wait_for_server() {
+  # wait for vllm server to start
+  # return 1 if vllm server crashes
+  local port=$1
+  timeout 1200 bash -c "
+    until curl -s localhost:${port}/v1/completions > /dev/null; do
+      sleep 1
+    done" && return 0 || return 1
+}
+
+
+benchmark() {
+
+  export VLLM_LOGGING_LEVEL=DEBUG
+  export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
+
+  # compare chunked prefill with disaggregated prefill
+
+  results_folder="./results"
+  model="meta-llama/Meta-Llama-3.1-8B-Instruct"
+  dataset_name="sonnet"
+  dataset_path="../sonnet_4x.txt"
+  num_prompts=10
+  qps=$1
+  prefix_len=50
+  input_len=2048
+  output_len=$2
+
+
+  CUDA_VISIBLE_DEVICES=0 vllm serve $model \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.6 \
+    --kv-transfer-config \
+    '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
+
+
+  CUDA_VISIBLE_DEVICES=1 vllm serve $model \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.6 \
+    --kv-transfer-config \
+    '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
+
+  wait_for_server 8100
+  wait_for_server 8200
+
+  # let the prefill instance finish prefill
+  vllm bench serve \
+    --backend vllm \
+    --model $model \
+    --dataset-name $dataset_name \
+    --dataset-path $dataset_path \
+    --sonnet-input-len $input_len \
+    --sonnet-output-len "$output_len" \
+    --sonnet-prefix-len $prefix_len \
+    --num-prompts $num_prompts \
+    --port 8100 \
+    --save-result \
+    --result-dir $results_folder \
+    --result-filename disagg_prefill_tp1.json \
+    --request-rate "inf"
+
+
+  # send the request to decode.
+  # The TTFT of this command will be the overhead of disagg prefill impl.
+  vllm bench serve \
+    --backend vllm \
+    --model $model \
+    --dataset-name $dataset_name \
+    --dataset-path $dataset_path \
+    --sonnet-input-len $input_len \
+    --sonnet-output-len "$output_len" \
+    --sonnet-prefix-len $prefix_len \
+    --num-prompts $num_prompts \
+    --port 8200 \
+    --save-result \
+    --result-dir $results_folder \
+    --result-filename disagg_prefill_tp1_overhead.json \
+    --request-rate "$qps"
+  kill_gpu_processes
+
+}
+
+
+main() {
+
+  (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
+  (which jq) || (apt-get -y install jq)
+  (which socat) || (apt-get -y install socat)
+
+  pip install quart httpx datasets
+
+  cd "$(dirname "$0")"
+
+  cd ..
+  # create sonnet-4x.txt
+  echo "" > sonnet_4x.txt
+  for _ in {1..4}
+  do
+    cat sonnet.txt >> sonnet_4x.txt
+  done
+  cd disagg_benchmarks
+
+  rm -rf results
+  mkdir results
+
+  default_qps=1
+  default_output_len=1
+  benchmark $default_qps $default_output_len
+
+}
+
+
+main "$@"
@@ -0,0 +1,157 @@
+#!/bin/bash
+
+# Requirement: 2x GPUs.
+
+
+# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
+# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
+# Resource: 2x GPU
+# Approaches:
+# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
+# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
+# Prefilling instance: max_output_token=1
+# Decoding instance: force the input tokens be the same across requests to bypass prefilling
+
+set -ex
+
+kill_gpu_processes() {
+  # kill all processes on GPU.
+  pgrep pt_main_thread | xargs -r kill -9
+  pgrep python3 | xargs -r kill -9
+  # vLLM now names the process with VLLM prefix after https://github.com/vllm-project/vllm/pull/21445
+  pgrep VLLM | xargs -r kill -9
+  for port in 8000 8100 8200; do lsof -t -i:$port | xargs -r kill -9; done
+  sleep 1
+}
+
+wait_for_server() {
+  # wait for vllm server to start
+  # return 1 if vllm server crashes
+  local port=$1
+  timeout 1200 bash -c "
+    until curl -s localhost:${port}/v1/completions > /dev/null; do
+      sleep 1
+    done" && return 0 || return 1
+}
+
+
+launch_chunked_prefill() {
+  model="meta-llama/Meta-Llama-3.1-8B-Instruct"
+  # disagg prefill
+  CUDA_VISIBLE_DEVICES=0 vllm serve $model \
+    --port 8100 \
+    --max-model-len 10000 \
+    --enable-chunked-prefill \
+    --gpu-memory-utilization 0.6 &
+  CUDA_VISIBLE_DEVICES=1 vllm serve $model \
+    --port 8200 \
+    --max-model-len 10000 \
+    --enable-chunked-prefill \
+    --gpu-memory-utilization 0.6 &
+  wait_for_server 8100
+  wait_for_server 8200
+  python3 round_robin_proxy.py &
+  sleep 1
+}
+
+
+launch_disagg_prefill() {
+  model="meta-llama/Meta-Llama-3.1-8B-Instruct"
+  # disagg prefill
+  CUDA_VISIBLE_DEVICES=0 vllm serve $model \
+    --port 8100 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.6 \
+    --kv-transfer-config \
+    '{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
+
+  CUDA_VISIBLE_DEVICES=1 vllm serve $model \
+    --port 8200 \
+    --max-model-len 10000 \
+    --gpu-memory-utilization 0.6 \
+    --kv-transfer-config \
+    '{"kv_connector":"P2pNcclConnector","kv_role":"kv_consumer","kv_rank":1,"kv_parallel_size":2,"kv_buffer_size":5e9}' &
+
+  wait_for_server 8100
+  wait_for_server 8200
+  python3 disagg_prefill_proxy_server.py &
+  sleep 1
+}
+
+
+benchmark() {
+  results_folder="./results"
+  model="meta-llama/Meta-Llama-3.1-8B-Instruct"
+  dataset_name="sonnet"
+  dataset_path="../sonnet_4x.txt"
+  num_prompts=100
+  qps=$1
+  prefix_len=50
+  input_len=1024
+  output_len=$2
+  tag=$3
+
+  vllm bench serve \
+    --backend vllm \
+    --model $model \
+    --dataset-name $dataset_name \
+    --dataset-path $dataset_path \
+    --sonnet-input-len $input_len \
+    --sonnet-output-len "$output_len" \
+    --sonnet-prefix-len $prefix_len \
+    --num-prompts $num_prompts \
+    --port 8000 \
+    --save-result \
+    --result-dir $results_folder \
+    --result-filename "$tag"-qps-"$qps".json \
+    --request-rate "$qps"
+
+  sleep 2
+}
+
+
+main() {
+
+  (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
+  (which jq) || (apt-get -y install jq)
+  (which socat) || (apt-get -y install socat)
+  (which lsof) || (apt-get -y install lsof)
+
+  pip install quart httpx matplotlib aiohttp datasets
+
+  cd "$(dirname "$0")"
+
+  cd ..
+  # create sonnet-4x.txt so that we can sample 2048 tokens for input
+  echo "" > sonnet_4x.txt
+  for _ in {1..4}
+  do
+    cat sonnet.txt >> sonnet_4x.txt
+  done
+  cd disagg_benchmarks
+
+  rm -rf results
+  mkdir results
+
+  default_output_len=6
+
+  export VLLM_HOST_IP=$(hostname -I | awk '{print $1}')
+
+  launch_chunked_prefill
+  for qps in 2 4 6 8; do
+  benchmark $qps $default_output_len chunked_prefill
+  done
+  kill_gpu_processes
+
+  launch_disagg_prefill
+  for qps in 2 4 6 8; do
+  benchmark $qps $default_output_len disagg_prefill
+  done
+  kill_gpu_processes
+
+  python3 visualize_benchmark_results.py
+
+}
+
+
+main "$@"
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import asyncio
+import logging
+import os
+import time
+import uuid
+from urllib.parse import urlparse
+
+import aiohttp
+from quart import Quart, Response, make_response, request
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    """parse command line arguments"""
+    parser = argparse.ArgumentParser(description="vLLM P/D disaggregation proxy server")
+
+    # Add args
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=6 * 60 * 60,
+        help="Timeout for backend service requests in seconds (default: 21600)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to run the server on (default: 8000)",
+    )
+    parser.add_argument(
+        "--prefill-url",
+        type=str,
+        default="http://localhost:8100",
+        help="Prefill service base URL (protocol + host[:port])",
+    )
+    parser.add_argument(
+        "--decode-url",
+        type=str,
+        default="http://localhost:8200",
+        help="Decode service base URL (protocol + host[:port])",
+    )
+    parser.add_argument(
+        "--kv-host",
+        type=str,
+        default="localhost",
+        help="Hostname or IP used by KV transfer (default: localhost)",
+    )
+    parser.add_argument(
+        "--prefill-kv-port",
+        type=int,
+        default=14579,
+        help="Prefill KV port (default: 14579)",
+    )
+    parser.add_argument(
+        "--decode-kv-port",
+        type=int,
+        default=14580,
+        help="Decode KV port (default: 14580)",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    """parse command line arguments"""
+    args = parse_args()
+
+    # Initialize configuration using command line parameters
+    AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
+    PREFILL_SERVICE_URL = args.prefill_url
+    DECODE_SERVICE_URL = args.decode_url
+    PORT = args.port
+
+    PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
+    DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
+
+    logger.info(
+        "Proxy resolved KV addresses -> prefill: %s, decode: %s",
+        PREFILL_KV_ADDR,
+        DECODE_KV_ADDR,
+    )
+
+    app = Quart(__name__)
+
+    # Attach the configuration object to the application instance so helper
+    # coroutines can read the resolved backend URLs and timeouts without using
+    # globals.
+    app.config.update(
+        {
+            "AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
+            "PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
+            "DECODE_SERVICE_URL": DECODE_SERVICE_URL,
+            "PREFILL_KV_ADDR": PREFILL_KV_ADDR,
+            "DECODE_KV_ADDR": DECODE_KV_ADDR,
+        }
+    )
+
+    def _normalize_base_url(url: str) -> str:
+        """Remove any trailing slash so path joins behave predictably."""
+        return url.rstrip("/")
+
+    def _get_host_port(url: str) -> str:
+        """Return the hostname:port portion for logging and KV headers."""
+        parsed = urlparse(url)
+        host = parsed.hostname or "localhost"
+        port = parsed.port
+        if port is None:
+            port = 80 if parsed.scheme == "http" else 443
+        return f"{host}:{port}"
+
+    PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
+    DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
+    KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
+
+    def _build_headers(request_id: str) -> dict[str, str]:
+        """Construct the headers expected by vLLM's P2P disagg connector."""
+        headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        return headers
+
+    async def _run_prefill(
+        request_path: str,
+        payload: dict,
+        headers: dict[str, str],
+        request_id: str,
+    ):
+        url = f"{PREFILL_BASE}{request_path}"
+        start_ts = time.perf_counter()
+        logger.info("[prefill] start request_id=%s url=%s", request_id, url)
+        try:
+            async with (
+                aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
+                session.post(url=url, json=payload, headers=headers) as resp,
+            ):
+                if resp.status != 200:
+                    error_text = await resp.text()
+                    raise RuntimeError(
+                        f"Prefill backend error {resp.status}: {error_text}"
+                    )
+                await resp.read()
+                logger.info(
+                    "[prefill] done request_id=%s status=%s elapsed=%.2fs",
+                    request_id,
+                    resp.status,
+                    time.perf_counter() - start_ts,
+                )
+        except asyncio.TimeoutError as exc:
+            raise RuntimeError(f"Prefill service timeout at {url}") from exc
+        except aiohttp.ClientError as exc:
+            raise RuntimeError(f"Prefill service unavailable at {url}") from exc
+
+    async def _stream_decode(
+        request_path: str,
+        payload: dict,
+        headers: dict[str, str],
+        request_id: str,
+    ):
+        url = f"{DECODE_BASE}{request_path}"
+        # Stream tokens from the decode service once the prefill stage has
+        # materialized KV caches on the target workers.
+        logger.info("[decode] start request_id=%s url=%s", request_id, url)
+        try:
+            async with (
+                aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
+                session.post(url=url, json=payload, headers=headers) as resp,
+            ):
+                if resp.status != 200:
+                    error_text = await resp.text()
+                    logger.error(
+                        "Decode backend error %s - %s", resp.status, error_text
+                    )
+                    err_msg = (
+                        '{"error": "Decode backend error ' + str(resp.status) + '"}'
+                    )
+                    yield err_msg.encode()
+                    return
+                logger.info(
+                    "[decode] streaming response request_id=%s status=%s",
+                    request_id,
+                    resp.status,
+                )
+                async for chunk_bytes in resp.content.iter_chunked(1024):
+                    yield chunk_bytes
+                logger.info("[decode] finished streaming request_id=%s", request_id)
+        except asyncio.TimeoutError:
+            logger.error("Decode service timeout at %s", url)
+            yield b'{"error": "Decode service timeout"}'
+        except aiohttp.ClientError as exc:
+            logger.error("Decode service error at %s: %s", url, exc)
+            yield b'{"error": "Decode service unavailable"}'
+
+    async def process_request():
+        """Process a single request through prefill and decode stages"""
+        try:
+            original_request_data = await request.get_json()
+
+            # Create prefill request (max_tokens=1)
+            prefill_request = original_request_data.copy()
+            prefill_request["max_tokens"] = 1
+            if "max_completion_tokens" in prefill_request:
+                prefill_request["max_completion_tokens"] = 1
+
+            # Execute prefill stage
+            # The request id encodes both KV socket addresses so the backend can
+            # shuttle tensors directly via NCCL once the prefill response
+            # completes.
+            request_id = (
+                f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
+                f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
+            )
+
+            headers = _build_headers(request_id)
+            await _run_prefill(request.path, prefill_request, headers, request_id)
+
+            # Execute decode stage and stream response
+            # Pass the unmodified user request so the decode phase can continue
+            # sampling with the already-populated KV cache.
+            generator = _stream_decode(
+                request.path, original_request_data, headers, request_id
+            )
+            response = await make_response(generator)
+            response.timeout = None  # Disable timeout for streaming response
+            return response
+
+        except Exception:
+            logger.exception("Error processing request")
+            return Response(
+                response=b'{"error": "Internal server error"}',
+                status=500,
+                content_type="application/json",
+            )
+
+    @app.route("/v1/completions", methods=["POST"])
+    async def handle_request():
+        """Handle incoming API requests with concurrency and rate limiting"""
+        try:
+            return await process_request()
+        except asyncio.CancelledError:
+            logger.warning("Request cancelled")
+            return Response(
+                response=b'{"error": "Request cancelled"}',
+                status=503,
+                content_type="application/json",
+            )
+
+    # Start the Quart server with host can be set to 0.0.0.0
+    app.run(port=PORT)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import itertools
+
+import aiohttp
+from aiohttp import web
+
+
+class RoundRobinProxy:
+    def __init__(self, target_ports):
+        self.target_ports = target_ports
+        self.port_cycle = itertools.cycle(self.target_ports)
+
+    async def handle_request(self, request):
+        target_port = next(self.port_cycle)
+        target_url = f"http://localhost:{target_port}{request.path_qs}"
+
+        async with aiohttp.ClientSession() as session:
+            try:
+                # Forward the request
+                async with session.request(
+                    method=request.method,
+                    url=target_url,
+                    headers=request.headers,
+                    data=request.content,
+                ) as response:
+                    # Start sending the response
+                    resp = web.StreamResponse(
+                        status=response.status, headers=response.headers
+                    )
+                    await resp.prepare(request)
+
+                    # Stream the response content
+                    async for chunk in response.content.iter_any():
+                        await resp.write(chunk)
+
+                    await resp.write_eof()
+                    return resp
+
+            except Exception as e:
+                return web.Response(text=f"Error: {str(e)}", status=500)
+
+
+async def main():
+    proxy = RoundRobinProxy([8100, 8200])
+    app = web.Application()
+    app.router.add_route("*", "/{path:.*}", proxy.handle_request)
+
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, "localhost", 8000)
+    await site.start()
+
+    print("Proxy server started on http://localhost:8000")
+
+    # Keep the server running
+    await asyncio.Event().wait()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import matplotlib.pyplot as plt
+import pandas as pd
+
+if __name__ == "__main__":
+    data = []
+    for name in ["disagg_prefill", "chunked_prefill"]:
+        for qps in [2, 4, 6, 8]:
+            with open(f"results/{name}-qps-{qps}.json") as f:
+                x = json.load(f)
+                x["name"] = name
+                x["qps"] = qps
+                data.append(x)
+
+    df = pd.DataFrame.from_dict(data)
+    dis_df = df[df["name"] == "disagg_prefill"]
+    chu_df = df[df["name"] == "chunked_prefill"]
+
+    plt.style.use("bmh")
+    plt.rcParams["font.size"] = 20
+
+    for key in [
+        "mean_ttft_ms",
+        "median_ttft_ms",
+        "p99_ttft_ms",
+        "mean_itl_ms",
+        "median_itl_ms",
+        "p99_itl_ms",
+    ]:
+        fig, ax = plt.subplots(figsize=(11, 7))
+        plt.plot(
+            dis_df["qps"], dis_df[key], label="disagg_prefill", marker="o", linewidth=4
+        )
+        plt.plot(
+            chu_df["qps"], chu_df[key], label="chunked_prefill", marker="o", linewidth=4
+        )
+        ax.legend()
+
+        ax.set_xlabel("QPS")
+        ax.set_ylabel(key)
+        ax.set_ylim(bottom=0)
+        fig.savefig(f"results/{key}.png")
+        plt.close(fig)
@@ -250,7 +250,7 @@ def benchmark_config(
                    num_experts=num_experts,
                    experts_per_token=topk,
                    hidden_dim=hidden_size,
-                    intermediate_size=shard_intermediate_size,
+                    intermediate_size_per_partition=shard_intermediate_size,
                    num_local_experts=num_experts,
                    num_logical_experts=num_experts,
                    activation=MoEActivation.SILU,
@@ -271,6 +271,7 @@ def benchmark_config(
                    moe_config=moe_config,
                    quant_config=quant_config,
                ),
+                inplace=not disable_inplace(),
            )

        with override_config(config):
@@ -278,6 +279,7 @@ def benchmark_config(
                x, input_gating, topk, renormalize=not use_deep_gemm
            )

+            inplace = not disable_inplace()
            if use_deep_gemm:
                return deep_gemm_experts.apply(
                    x,
@@ -296,6 +298,7 @@ def benchmark_config(
                w2,
                topk_weights,
                topk_ids,
+                inplace=inplace,
                quant_config=quant_config,
            )

@@ -65,32 +65,6 @@ class RequestArgs(NamedTuple):
    limit_min_tokens: int  # Use negative value for no limit
    limit_max_tokens: int  # Use negative value for no limit
    timeout_sec: int
-    send_conversation_id: bool
-    headers: dict[str, str]
-
-
-def parse_custom_header(header: str) -> tuple[str, str]:
-    separators = (":", "=")
-    for separator in separators:
-        if separator in header:
-            key, value = header.split(separator, 1)
-            key = key.strip()
-            value = value.strip()
-            if key:
-                return key, value
-            break
-    raise argparse.ArgumentTypeError(
-        "Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
-    )
-
-
-def build_request_headers(
-    api_key: str | None, custom_headers: list[tuple[str, str]] | None
-) -> dict[str, str]:
-    headers = dict(custom_headers or [])
-    if api_key:
-        headers["Authorization"] = f"Bearer {api_key}"
-    return headers


 class BenchmarkArgs(NamedTuple):
@@ -244,11 +218,12 @@ async def send_request(
    max_tokens: int | None = None,
    timeout_sec: int = 120,
    conversation_id: str | None = None,
-    headers: dict[str, str] | None = None,
 ) -> ServerResponse:
    payload = {
        "model": model,
        "messages": messages,
+        "seed": 0,
+        "temperature": 0.0,
    }

    if conversation_id is not None:
@@ -258,17 +233,13 @@ async def send_request(
        payload["stream"] = True
        payload["stream_options"] = {"include_usage": False}

-    # if min_tokens is not None:
-    #     payload["min_tokens"] = min_tokens
+    if min_tokens is not None:
+        payload["min_tokens"] = min_tokens

    if max_tokens is not None:
        payload["max_tokens"] = max_tokens

-    request_headers = {"Content-Type": "application/json"}
-    if conversation_id is not None:
-        request_headers["X-Session-ID"] = str(conversation_id)
-    if headers is not None:
-        request_headers.update(headers)
+    headers = {"Content-Type": "application/json"}

    # Calculate the timeout for the request
    if max_tokens is not None:
@@ -294,7 +265,7 @@ async def send_request(
    most_recent_timestamp: int = start_time

    async with session.post(
-        url=chat_url, json=payload, headers=request_headers, timeout=timeout
+        url=chat_url, json=payload, headers=headers, timeout=timeout
    ) as response:
        http_status = HTTPStatus(response.status)
        if http_status == HTTPStatus.OK:
@@ -346,8 +317,6 @@ async def send_request(
            latency = time.perf_counter_ns() - start_time

    if ttft is None:
-        if stream:
-            valid_response = False
        # The response was a single chunk
        ttft = latency

@@ -454,8 +423,7 @@ async def send_turn(
        min_tokens,
        max_tokens,
        req_args.timeout_sec,
-        conversation_id=conv_id if req_args.send_conversation_id else None,
-        headers=req_args.headers,
+        conversation_id=conv_id,
    )

    if response.valid is False:
@@ -904,7 +872,6 @@ def get_client_config(
    # Arguments for API requests
    chat_url = f"{args.url}/v1/chat/completions"
    model_name = args.served_model_name if args.served_model_name else args.model
-    headers = build_request_headers(args.api_key, args.header)

    req_args = RequestArgs(
        chat_url=chat_url,
@@ -913,8 +880,6 @@ def get_client_config(
        limit_min_tokens=args.limit_min_tokens,
        limit_max_tokens=args.limit_max_tokens,
        timeout_sec=args.request_timeout_sec,
-        send_conversation_id=args.send_conversation_id,
-        headers=headers,
    )

    return client_args, req_args
@@ -1280,19 +1245,19 @@ def process_statistics(
        )


-async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
+async def get_server_info(url: str) -> None:
    logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
    async with aiohttp.ClientSession() as session:
        # Get server version (not mandatory, "version" endpoint may not exist)
        url_version = f"{url}/version"
-        async with session.get(url_version, headers=headers) as response:
+        async with session.get(url_version) as response:
            if HTTPStatus(response.status) == HTTPStatus.OK:
                text = await response.text()
                logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")

        # Get available models
        url_models = f"{url}/v1/models"
-        async with session.get(url_models, headers=headers) as response:
+        async with session.get(url_models) as response:
            if HTTPStatus(response.status) == HTTPStatus.OK:
                text = await response.text()
                logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1358,22 +1323,6 @@ async def main() -> None:
        help="Base URL for the LLM API server",
    )

-    parser.add_argument(
-        "--api-key",
-        type=str,
-        default=None,
-        help="API key to send as an Authorization bearer token",
-    )
-    parser.add_argument(
-        "--header",
-        action="append",
-        type=parse_custom_header,
-        default=None,
-        metavar="KEY=VALUE",
-        help="Custom request header. Can be specified multiple times. "
-        "Accepts 'Header-Name: value' or 'Header-Name=value'.",
-    )
-
    parser.add_argument(
        "-p",
        "--num-clients",
@@ -1488,22 +1437,6 @@ async def main() -> None:
        help="Disable stream/streaming mode (set 'stream' to False in the API request)",
    )

-    parser.add_argument(
-        "--send-conversation-id",
-        default=False,
-        action="store_true",
-        help=(
-            "Inject a `conversation_id` field into each Chat Completions "
-            "payload. This is a non-standard OpenAI extension consumed by "
-            "vLLM's disaggregated multi-turn proxy "
-            "(examples/disaggregated/disaggregated_serving/"
-            "disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. "
-            "Leave disabled (default) when targeting strict "
-            "OpenAI-compatible endpoints; enable when benchmarking the "
-            "disaggregated proxy."
-        ),
-    )
-
    parser.add_argument(
        "-e",
        "--excel-output",
@@ -1592,8 +1525,7 @@ async def main() -> None:
        args.model, trust_remote_code=args.trust_remote_code
    )

-    headers = build_request_headers(args.api_key, args.header)
-    await get_server_info(args.url, headers=headers)
+    await get_server_info(args.url)

    # Load the input file (either conversations of configuration file)
    logger.info(f"Reading input file: {args.input_file}")
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Build vLLM Rust artifacts and install them into the vllm package.
+# Build the vllm-rs Rust frontend binary and install it into the vllm package.
 # Usage: ./build_rust.sh [--debug]
 #
 # By default builds in release mode. Pass --debug for faster compile times
@@ -8,6 +8,8 @@
 set -euo pipefail

 REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
+RUST_DIR="$REPO_ROOT/rust"
+TARGET_PATH="${VLLM_RS_TARGET_PATH:-$REPO_ROOT/vllm/vllm-rs}"

 # Read the required toolchain from rust-toolchain.toml.
 TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -25,9 +27,18 @@ if ! rustup run "$TOOLCHAIN" rustc --version &>/dev/null; then
 fi

 if [[ "${1:-}" == "--debug" ]]; then
-    PROFILE_ARG="--debug"
+    PROFILE_ARGS=()
+    PROFILE_DIR="debug"
 else
-    PROFILE_ARG="--release"
+    PROFILE_ARGS=(--release)
+    PROFILE_DIR="release"
 fi

-python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
+cargo +"$TOOLCHAIN" build "${PROFILE_ARGS[@]}" \
+    --manifest-path "$RUST_DIR/Cargo.toml" \
+    --bin vllm-rs \
+    --features native-tls-vendored
+
+mkdir -p "$(dirname "$TARGET_PATH")"
+cp "$RUST_DIR/target/$PROFILE_DIR/vllm-rs" "$TARGET_PATH"
+echo "Installed vllm-rs to $TARGET_PATH"
@@ -14,18 +14,7 @@ import argparse
 import os
 import shutil

-from torch.utils.hipify.hipify_python import get_hip_file_path, hipify
-
-
-def _expected_hip_build_path(source_abs: str, output_directory: str) -> str:
-    """Match torch.utils.hipify.hipify_python.preprocessor fout_path naming."""
-    rel = os.path.relpath(source_abs, output_directory)
-    return os.path.abspath(
-        os.path.join(
-            output_directory, get_hip_file_path(rel, is_pytorch_extension=True)
-        )
-    )
-
+from torch.utils.hipify.hipify_python import hipify

 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
@@ -64,11 +53,7 @@ if __name__ == "__main__":
    hipify_result = hipify(
        project_directory=args.project_dir,
        output_directory=args.output_dir,
-        # Hipify resolves quoted includes next to the including file first; vLLM
-        # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h"
-        # from quantization/w8a8/fp8/*.cu). Without an include root here, those
-        # headers are never found and are not hipified or rewritten in dependents.
-        header_include_dirs=["."],
+        header_include_dirs=[],
        includes=includes,
        extra_files=extra_files,
        show_detailed=True,
@@ -79,20 +64,14 @@ if __name__ == "__main__":
    hipified_sources = []
    for source in args.sources:
        s_abs = os.path.abspath(source)
-        if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:
-            path = hipify_result[s_abs].hipified_path
-            # PyTorch skips writing when is_pytorch_extension and text unchanged;
-            # hipified_path then stays *.cu. CMake expects *.hip under output_dir.
-            if s_abs.endswith(".cu") and path.endswith(".cu"):
-                dest = _expected_hip_build_path(s_abs, args.output_dir)
-                if os.path.normpath(path) != os.path.normpath(dest):
-                    os.makedirs(os.path.dirname(dest), exist_ok=True)
-                    shutil.copy2(path, dest)
-                hipified_s_abs = dest
-            else:
-                hipified_s_abs = path
-        else:
-            hipified_s_abs = s_abs
+        hipified_s_abs = (
+            hipify_result[s_abs].hipified_path
+            if (
+                s_abs in hipify_result
+                and hipify_result[s_abs].hipified_path is not None
+            )
+            else s_abs
+        )
        hipified_sources.append(hipified_s_abs)

    assert len(hipified_sources) == len(args.sources)
@@ -81,14 +81,6 @@ function (hipify_sources_target OUT_SRCS NAME ORIG_SRCS)
  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_SRCS ${SRCS})
  set_property(GLOBAL APPEND PROPERTY VLLM_HIPIFY_ALL_BYPRODUCTS ${HIP_SRCS})

-  # Chain hipify targets so they run sequentially. Parallel hipify
-  # invocations race on shutil.copytree, overwriting .hip files
-  # produced by another target back to .cu originals.
-  if (DEFINED _VLLM_LAST_HIPIFY_TARGET)
-    add_dependencies(hipify${NAME} ${_VLLM_LAST_HIPIFY_TARGET})
-  endif()
-  set(_VLLM_LAST_HIPIFY_TARGET "hipify${NAME}" PARENT_SCOPE)
-
  # Swap out original extension sources with hipified sources.
  list(APPEND HIP_SRCS ${CXX_SRCS})
  set(${OUT_SRCS} ${HIP_SRCS} PARENT_SCOPE)
@@ -30,12 +30,7 @@
  }()

 namespace {
-enum class FusedMOEAct {
-  SiluAndMul,
-  SwigluOAIAndMul,
-  GeluAndMul,
-  GeluTanhAndMul,
-};
+enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };

 FusedMOEAct get_act_type(const std::string& act) {
  if (act == "silu") {
@@ -44,8 +39,6 @@ FusedMOEAct get_act_type(const std::string& act) {
    return FusedMOEAct::SwigluOAIAndMul;
  } else if (act == "gelu") {
    return FusedMOEAct::GeluAndMul;
-  } else if (act == "gelu_tanh") {
-    return FusedMOEAct::GeluTanhAndMul;
  } else {
    TORCH_CHECK(false, "Invalid act type: " + act);
  }
@@ -150,44 +143,6 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
  }
 }

-template <typename scalar_t>
-void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
-                       const int32_t m_size, const int32_t n_size,
-                       const int32_t input_stride,
-                       const int32_t output_stride) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  const int32_t dim = n_size / 2;
-  float* __restrict__ gate = input;
-  float* __restrict__ up = input + dim;
-  vec_op::FP32Vec16 one_vec(1.0);
-  vec_op::FP32Vec16 w1_vec(0.7978845608028654);
-  vec_op::FP32Vec16 w2_vec(0.5);
-  vec_op::FP32Vec16 w3_vec(0.044715);
-  alignas(64) float temp[16];
-
-  for (int32_t m = 0; m < m_size; ++m) {
-    for (int32_t n = 0; n < dim; n += 16) {
-      vec_op::FP32Vec16 gate_vec(gate + n);
-      vec_op::FP32Vec16 up_vec(up + n);
-      auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
-      auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
-
-      inner_vec.save(temp);
-      for (int32_t i = 0; i < 16; ++i) {
-        temp[i] = std::tanh(temp[i]);
-      }
-      vec_op::FP32Vec16 tanh_vec(temp);
-      auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
-      auto gated_output_fp32 = up_vec * gelu_tanh;
-      scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
-      gated_output.save(output + n);
-    }
-    gate += input_stride;
-    up += input_stride;
-    output += output_stride;
-  }
-}
-
 template <typename scalar_t>
 FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
                                  float* __restrict__ input,
@@ -205,9 +160,6 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
    case FusedMOEAct::GeluAndMul:
      gelu_and_mul(input, output, m, n, input_stride, output_stride);
      return;
-    case FusedMOEAct::GeluTanhAndMul:
-      gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
-      return;
    default:
      TORCH_CHECK(false, "Unsupported act type.");
  }
@@ -1,4 +1,3 @@
-// TODO: Remove this once ROCm upgrade to torch 2.11.
 #include <torch/all.h>
 #include <torch/cuda.h>
 #include <cuda_runtime.h>
@@ -9,7 +9,6 @@
 static const char* PYARGS_PARSE = "KKKK";
 #else
  #include <cstdlib>
-  #include <cstdint>
  #include <cerrno>
  #include <climits>

@@ -47,29 +46,6 @@ static inline unsigned long long my_min(unsigned long long a,
  return a < b ? a : b;
 }

-static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,
-                                     size_t alignment) {
-  CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);
-  if (status == CUresult(0) || alignment == 0) {
-    return status;
-  }
-
-  // Some ROCm stacks can report OOM while reserving VA with an explicit
-  // alignment even when physical VRAM is free. Let HIP choose the default
-  // alignment, then verify that the returned address still satisfies the
-  // requested alignment before accepting it.
-  status = cuMemAddressReserve(d_mem, size, 0, 0, 0);
-  if (status != CUresult(0)) {
-    return status;
-  }
-  if (((std::uintptr_t)(*d_mem) % alignment) == 0) {
-    return status;
-  }
-
-  (void)cuMemAddressFree(*d_mem, size);
-  return hipErrorNotSupported;
-}
-
 static const char* PYARGS_PARSE = "KKKO";
 #endif

@@ -349,7 +325,7 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
    return nullptr;
  }
 #else
-  CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity));
+  CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
  if (error_code != 0) {
    return nullptr;
  }
@@ -1,11 +1,7 @@
-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>

 #include "custom_all_reduce.cuh"

@@ -15,7 +11,7 @@ using fptr_t = int64_t;
 static_assert(sizeof(void*) == sizeof(fptr_t));

 fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
-                      torch::stable::Tensor& rank_data, int64_t rank,
+                      torch::Tensor& rank_data, int64_t rank,
                      bool fully_connected) {
  int world_size = fake_ipc_ptrs.size();
  if (world_size > 8)
@@ -29,9 +25,9 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
  for (int i = 0; i < world_size; i++) {
    ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
  }
-  return (fptr_t) new vllm::CustomAllreduce(
-      ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank,
-      world_size, fully_connected);
+  return (fptr_t) new vllm::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(),
+                                            rank_data.numel(), rank, world_size,
+                                            fully_connected);
 }

 /**
@@ -50,14 +46,10 @@ fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
 * 5. A[None].expand(2, -1, -1, -1): Not OK
 * 6. A[:, 1:, 1:]: Not OK
 */
-bool _is_weak_contiguous(torch::stable::Tensor& t) {
-  if (t.is_contiguous()) {
-    return true;
-  }
-  int64_t storage_nbytes = 0;
-  TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes));
-  return storage_nbytes - t.storage_offset() * t.element_size() ==
-         static_cast<int64_t>(t.numel() * t.element_size());
+bool _is_weak_contiguous(torch::Tensor& t) {
+  return t.is_contiguous() ||
+         (t.storage().nbytes() - t.storage_offset() * t.element_size() ==
+          t.numel() * t.element_size());
 }

 /**
@@ -67,45 +59,42 @@ bool _is_weak_contiguous(torch::stable::Tensor& t) {
 * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
 * copied into _reg_buffer.
 */
-void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
-                torch::stable::Tensor& out, fptr_t _reg_buffer,
-                int64_t reg_buffer_sz_bytes) {
+void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
+                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      inp.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();

-  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
-  STD_TORCH_CHECK((inp.numel()) == (out.numel()));
-  STD_TORCH_CHECK(_is_weak_contiguous(out));
-  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
+  TORCH_CHECK_EQ(inp.numel(), out.numel());
+  TORCH_CHECK(_is_weak_contiguous(out));
+  TORCH_CHECK(_is_weak_contiguous(inp));
  auto input_size = inp.numel() * inp.element_size();
  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
  if (reg_buffer) {
-    STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
-    STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
-                                   cudaMemcpyDeviceToDevice, stream));
+    TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
+    AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
+                                  cudaMemcpyDeviceToDevice, stream));
  } else {
-    reg_buffer = inp.mutable_data_ptr();
+    reg_buffer = inp.data_ptr();
  }
  switch (out.scalar_type()) {
-    case torch::headeronly::ScalarType::Float: {
+    case at::ScalarType::Float: {
      fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
-                           reinterpret_cast<float*>(out.mutable_data_ptr()),
+                           reinterpret_cast<float*>(out.data_ptr()),
                           out.numel());
      break;
    }
-    case torch::headeronly::ScalarType::Half: {
+    case at::ScalarType::Half: {
      fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
-                          reinterpret_cast<half*>(out.mutable_data_ptr()),
-                          out.numel());
+                          reinterpret_cast<half*>(out.data_ptr()), out.numel());
      break;
    }
 #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-    case torch::headeronly::ScalarType::BFloat16: {
+    case at::ScalarType::BFloat16: {
      fa->allreduce<nv_bfloat16>(
          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
-          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), out.numel());
+          reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
      break;
    }
 #endif
@@ -123,7 +112,7 @@ int64_t meta_size() { return sizeof(vllm::Signal); }

 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
+  TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
  void* ipc_ptrs[8];
  for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
    ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
@@ -154,49 +143,47 @@ void register_graph_buffers(fptr_t _fa,
  fa->register_graph_buffers(bytes, offsets);
 }

-std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
+std::tuple<fptr_t, torch::Tensor> allocate_shared_buffer_and_handle(
    int64_t size) {
-  int device_index;
-  STD_CUDA_CHECK(cudaGetDevice(&device_index));
-  const torch::stable::accelerator::DeviceGuard device_guard(device_index);
+  auto device_index = c10::cuda::current_device();
+  at::DeviceGuard device_guard(at::Device(at::DeviceType::CUDA, device_index));
  void* buffer;
  cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
-  const cudaStream_t stream = get_current_cuda_stream(device_index);
-  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));

  // Allocate buffer
 #if defined(USE_ROCM)
  // data buffers need to be "uncached" for signal on MI200
-  STD_CUDA_CHECK(
+  AT_CUDA_CHECK(
      hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
 #else
-  STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
+  AT_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
 #endif
-  STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
-  STD_CUDA_CHECK(cudaStreamSynchronize(stream));
-  STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
+  AT_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
+  AT_CUDA_CHECK(cudaStreamSynchronize(stream));
+  AT_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));

  // Create IPC memhandle for the allocated buffer.
  // Will use it in open_mem_handle.
-  auto handle = torch::stable::empty(
-      {static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))},
-      torch::headeronly::ScalarType::Byte, std::nullopt,
-      torch::stable::Device(torch::stable::DeviceType::CPU));
-  STD_CUDA_CHECK(cudaIpcGetMemHandle(
-      (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer));
+  auto options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
+  auto handle =
+      torch::empty({static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))}, options);
+  AT_CUDA_CHECK(
+      cudaIpcGetMemHandle((cudaIpcMemHandle_t*)handle.data_ptr(), buffer));

  return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
 }

-fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
+fptr_t open_mem_handle(torch::Tensor& mem_handle) {
  void* ipc_ptr;
-  STD_CUDA_CHECK(cudaIpcOpenMemHandle(
-      (void**)&ipc_ptr,
-      *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()),
+  AT_CUDA_CHECK(cudaIpcOpenMemHandle(
+      (void**)&ipc_ptr, *((const cudaIpcMemHandle_t*)mem_handle.data_ptr()),
      cudaIpcMemLazyEnablePeerAccess));
  return reinterpret_cast<fptr_t>(ipc_ptr);
 }

 void free_shared_buffer(fptr_t buffer) {
-  STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
+  AT_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
 }
@@ -1,4 +1,4 @@
-#include "common.hpp"
+#include "cutlass_extensions/common.hpp"

 int32_t get_sm_version_num() {
  int32_t major_capability, minor_capability;
@@ -28,20 +28,7 @@
 *   [bs*576,       bs*576 + bs*8):   UE8M0 scales, 7 real + 1 pad per token
 */

-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-
 #include <cmath>
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-#include "type_convert.cuh"
-
 #ifndef USE_ROCM
  #include <cuda_fp8.h>
 #else
@@ -50,6 +37,14 @@
 #include <cuda_runtime.h>
 #include <type_traits>

+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/cuda.h>
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+#include "type_convert.cuh"
+
 #ifndef FINAL_MASK
  #ifdef USE_ROCM
    #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,7 +70,7 @@ namespace deepseek_v4_fused_ops {

 namespace {
 inline int getSMVersion() {
-  auto* props = get_device_prop();
+  auto* props = at::cuda::getCurrentDeviceProperties();
  return props->major * 10 + props->minor;
 }
 }  // namespace
@@ -102,35 +97,6 @@ constexpr float NUM_TOKEN_CUTOFF = 1024;
 constexpr int kNumLanes = 32;
 constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 16

-// Pack this lane's 16 fp32 elements into per-tensor E4M3 FP8 (one uint4 = 16
-// B), scaling by `scale` (a reciprocal scale) and saturating to ±448.  Used by
-// the FlashInfer full-cache path for both the Q and KV stores.
-__device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
-                                                float const scale) {
-#ifndef USE_ROCM
-  uint4 out;
-  auto* out2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane / 2; i++) {
-    float2 scaled =
-        make_float2(values[2 * i] * scale, values[2 * i + 1] * scale);
-    scaled.x = fminf(fmaxf(scaled.x, -kFp8Max), kFp8Max);
-    scaled.y = fminf(fmaxf(scaled.y, -kFp8Max), kFp8Max);
-    out2[i] = __nv_cvt_float2_to_fp8x2(scaled, __NV_SATFINITE, __NV_E4M3);
-  }
-  return out;
-#else
-  uint8_t out_bytes[kElemsPerLane];
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane; i++) {
-    float scaled = values[i] * scale;
-    scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
-    out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
-  }
-  return *reinterpret_cast<uint4 const*>(out_bytes);
-#endif
-}
-
 // ────────────────────────────────────────────────────────────────────────────
 // Small inline helpers
 // ────────────────────────────────────────────────────────────────────────────
@@ -598,7 +564,7 @@ static void launchFusedDeepseekV4Templated(
  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
  // unavailable there.  Refuse the launch loudly instead of silently
  // skipping the work.
-  STD_TORCH_CHECK(
+  TORCH_CHECK(
      sm_version >= 80,
      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
      "(Ampere or newer); got sm_",
@@ -669,7 +635,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
    DISPATCH(64)
    DISPATCH(128)
    default:
-      STD_TORCH_CHECK(false,
+      TORCH_CHECK(false,
                  "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
                  "unsupported num_heads_q_padded=",
                  num_heads_q_padded,
@@ -678,504 +644,80 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 #undef DISPATCH
 }

-// ────────────────────────────────────────────────────────────────────────────
-// FlashInfer full-cache kernel
-// ────────────────────────────────────────────────────────────────────────────
-//
-// Sibling to the FlashMLA kernel above, used by the FlashInfer V4 sparse-MLA
-// backend.  Differences from the legacy path:
-//   * No Q head padding — output Q layout matches the input num_heads_q.
-//   * KV is written as a *contiguous* 512-wide row per token (token-strided),
-//     not the legacy UE8M0 paged layout with a separate scale tail.
-//   * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
-//     scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
-//
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
-// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
-// num_heads_q → KV.
-template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
-__global__ void fusedDeepseekV4FullCacheKernel(
-    scalar_t_in* __restrict__ q_inout,          // [N, H, 512], in place (bf16)
-    uint8_t* __restrict__ q_fp8_out,            // [N, H, 512] fp8, optional
-    int64_t const q_fp8_stride0,                // elements (fp8 == bytes)
-    int64_t const q_fp8_stride1,                // elements (fp8 == bytes)
-    scalar_t_in const* __restrict__ kv_in,      // [N, 512] bf16
-    uint8_t* __restrict__ k_cache,              // contiguous bf16 or fp8 cache
-    int64_t const* __restrict__ slot_mapping,   // [num_tokens_insert] i64
-    int64_t const* __restrict__ position_ids,   // [N] i64
-    float const* __restrict__ cos_sin_cache,    // [max_pos, 64] fp32
-    float const* __restrict__ fp8_scale_ptr,    // scalar, KV fp8 only
-    float const* __restrict__ q_fp8_scale_inv,  // scalar, Q fp8 only
-    float const eps,
-    int const num_tokens_full,      // = q.size(0) = kv.size(0)
-    int const num_tokens_insert,    // = slot_mapping.size(0)
-    int const num_heads_q,          // H (no padding)
-    int const cache_block_size,     // tokens per cache block
-    int64_t const kv_block_stride,  // bytes per cache block
-    int64_t const kv_token_stride) {  // bytes per cache token
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
-    return;
-  } else {
-#endif
-    using Converter = vllm::_typeConvert<scalar_t_in>;
-    int const warpsPerBlock = blockDim.x / 32;
-    int const warpId = threadIdx.x / 32;
-    int const laneId = threadIdx.x % 32;
-    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
-
-    int const slotsPerToken = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / slotsPerToken;
-    int const slotIdx = globalWarpIdx % slotsPerToken;
-    if (tokenIdx >= num_tokens_full) return;
-    bool const isKV = (slotIdx == num_heads_q);
-    // KV branch: skip DP-padded tokens (no slot reserved for them).
-    if (isKV && tokenIdx >= num_tokens_insert) return;
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaGridDependencySynchronize();
-#endif
-
-    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      src_ptr = q_inout +
-                (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
-                    kHeadDim +
-                dim_base;
-    }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
-
-    // ── Decode bf16 → 16 fp32 registers ───────────────────────────────────
-    float elements[kElemsPerLane];
-    {
-      auto const* p0 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
-      auto const* p1 =
-          reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p0[i]);
-        elements[2 * i] = f2.x;
-        elements[2 * i + 1] = f2.y;
-      }
-#pragma unroll
-      for (int i = 0; i < 4; i++) {
-        float2 f2 = Converter::convert(p1[i]);
-        elements[8 + 2 * i] = f2.x;
-        elements[8 + 2 * i + 1] = f2.y;
-      }
-    }
-
-    // ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
-    if (!isKV) {
-      float sumOfSquares = 0.0f;
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        sumOfSquares += elements[i] * elements[i];
-      }
-      sumOfSquares = warpSum<float>(sumOfSquares);
-      float const rms_rcp =
-          rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        elements[i] = elements[i] * rms_rcp;
-      }
-    }
-
-    // ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────
-    bool const is_rope_lane = dim_base >= kNopeDim;
-    if (is_rope_lane) {
-      int64_t const pos = position_ids[tokenIdx];
-      constexpr int kHalfRope = kRopeDim / 2;
-      float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
-      float const* sin_ptr = cos_ptr + kHalfRope;
-      int const rope_local_base = dim_base - kNopeDim;
-      int const half_base = rope_local_base >> 1;
-      float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
-      float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
-      float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
-      float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
-      float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
-      float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
-#pragma unroll
-      for (int p = 0; p < kElemsPerLane / 2; p++) {
-        float const x_even = elements[2 * p];
-        float const x_odd = elements[2 * p + 1];
-        elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
-        elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
-      }
-    }
-
-    // ── Store ─────────────────────────────────────────────────────────────
-    if (!isKV) {
-      if constexpr (STORE_Q_FP8) {
-        float const scale_inv = VLLM_LDG(q_fp8_scale_inv);
-        uint4 const out = packFp8E4M3x16(elements, scale_inv);
-        uint8_t* dst = q_fp8_out +
-                       static_cast<int64_t>(tokenIdx) * q_fp8_stride0 +
-                       static_cast<int64_t>(slotIdx) * q_fp8_stride1 + dim_base;
-        *reinterpret_cast<uint4*>(dst) = out;
-      } else {
-        uint4 out0, out1;
-        auto* po0 = reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-        auto* po1 = reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          po0[i] = Converter::convert(
-              make_float2(elements[2 * i], elements[2 * i + 1]));
-        }
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          po1[i] = Converter::convert(
-              make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-        }
-        scalar_t_in* dst =
-            q_inout +
-            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-            dim_base;
-        *reinterpret_cast<uint4*>(dst) = out0;
-        *reinterpret_cast<uint4*>(dst + 8) = out1;
-      }
-    } else {
-      int64_t const slot_id = slot_mapping[tokenIdx];
-      if (slot_id >= 0) {
-        int64_t const block_idx = slot_id / cache_block_size;
-        int64_t const pos_in_block = slot_id % cache_block_size;
-        uint8_t* cache_row =
-            k_cache + block_idx * kv_block_stride + pos_in_block * kv_token_stride;
-        if constexpr (STORE_KV_FP8) {
-          float const inv_scale = 1.0f / VLLM_LDG(fp8_scale_ptr);
-          uint4 const out = packFp8E4M3x16(elements, inv_scale);
-          *reinterpret_cast<uint4*>(cache_row + dim_base) = out;
-        } else {
-          uint4 out0, out1;
-          auto* po0 =
-              reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
-          auto* po1 =
-              reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
-#pragma unroll
-          for (int i = 0; i < 4; i++) {
-            po0[i] = Converter::convert(
-                make_float2(elements[2 * i], elements[2 * i + 1]));
-          }
-#pragma unroll
-          for (int i = 0; i < 4; i++) {
-            po1[i] = Converter::convert(
-                make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
-          }
-          scalar_t_in* dst = reinterpret_cast<scalar_t_in*>(cache_row) + dim_base;
-          *reinterpret_cast<uint4*>(dst) = out0;
-          *reinterpret_cast<uint4*>(dst + 8) = out1;
-        }
-      }
-    }
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaTriggerProgrammaticLaunchCompletion();
-#endif
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  }
-#endif
-}
-
-// Configure + launch helper shared by the bf16 and fp8 full-cache launchers.
-template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
-static void launchFullCacheKernel(
-    scalar_t_in* q_inout, uint8_t* q_fp8_out, int64_t q_fp8_stride0,
-    int64_t q_fp8_stride1, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
-    float const* cos_sin_cache, float const* fp8_scale,
-    float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
-    int const num_tokens_insert, int const num_heads_q,
-    int const cache_block_size, int64_t const kv_block_stride,
-    int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
-  constexpr int kBlockSize = 256;
-  constexpr int kWarpsPerBlock = kBlockSize / 32;
-  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
-  int const grid =
-      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
-  auto* kernel =
-      fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
-#ifndef USE_ROCM
-  static int const sm_version = getSMVersion();
-  STD_TORCH_CHECK(sm_version >= 80, op_name,
-                  " requires sm_80+ (Ampere or newer); got sm_", sm_version);
-  cudaLaunchConfig_t config;
-  config.gridDim = dim3(grid);
-  config.blockDim = dim3(kBlockSize);
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.attrs = attrs;
-  config.numAttrs = (sm_version >= 90) ? 1 : 0;
-  cudaLaunchKernelEx(&config, kernel, q_inout, q_fp8_out, q_fp8_stride0,
-                     q_fp8_stride1, kv_in, k_cache, slot_mapping, position_ids,
-                     cos_sin_cache, fp8_scale, q_fp8_scale_inv, eps,
-                     num_tokens_full, num_tokens_insert, num_heads_q,
-                     cache_block_size, kv_block_stride, kv_token_stride);
-#else
-  kernel<<<grid, kBlockSize, 0, stream>>>(
-      q_inout, q_fp8_out, q_fp8_stride0, q_fp8_stride1, kv_in, k_cache,
-      slot_mapping, position_ids, cos_sin_cache, fp8_scale, q_fp8_scale_inv,
-      eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
-      kv_block_stride, kv_token_stride);
-#endif
-}
-
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm

 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
-    torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
-    torch::stable::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
-    torch::stable::Tensor const& slot_mapping,   // [N] int64
-    torch::stable::Tensor const& position_ids,   // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
-    int64_t q_head_padded,                       // padded Q head count for output
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
+    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
+    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
+    torch::Tensor const& slot_mapping,   // [N] int64
+    torch::Tensor const& position_ids,   // [N] int64
+    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    int64_t q_head_padded,               // padded Q head count for output
    double eps, int64_t cache_block_size) {
-  STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(),
-                  "q_in must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() ==
-                          torch::headeronly::ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() ==
-                          torch::headeronly::ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda(), "cos_sin_cache must be CUDA");
-  STD_TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
-                  "q_in shape [N, num_heads_q, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
-                  "q_in and kv dtype must match");
-  STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
-                  "q_head_padded must be >= q_in.size(1) (num_heads_q)");
-  STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
-                  "k_cache must be uint8");
-  STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64]");
-  STD_TORCH_CHECK(cos_sin_cache.scalar_type() ==
-                      torch::headeronly::ScalarType::Float,
-                  "cos_sin_cache must be float32");
+  TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
+              "q_in must be contiguous CUDA");
+  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
+  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
+  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
+              "slot_mapping must be int64 CUDA");
+  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
+              "position_ids must be int64 CUDA");
+  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
+  TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
+              "q_in shape [N, num_heads_q, 512]");
+  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
+  TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
+  TORCH_CHECK(q_head_padded >= q_in.size(1),
+              "q_head_padded must be >= q_in.size(1) (num_heads_q)");
+  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
+  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
+              "cos_sin_cache shape [max_pos, 64]");
+  TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
+              "cos_sin_cache must be float32");

  // With DP padding, slot_mapping can be shorter than q/kv/positions.
  // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
  // KV quant+insert runs only on the first slot_mapping.size(0) rows.
  int const num_tokens_full = static_cast<int>(q_in.size(0));
  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
+  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
+                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
+              "q/kv/position_ids row counts must match");
+  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
+              "slot_mapping must not exceed q row count");
  int const num_heads_q = static_cast<int>(q_in.size(1));
  int const num_heads_q_padded = static_cast<int>(q_head_padded);
  int const cache_block_size_i = static_cast<int>(cache_block_size);
  int const kv_block_stride = static_cast<int>(k_cache.stride(0));

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q_in.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
+  auto stream = at::cuda::getCurrentCUDAStream();

  // Allocate the padded q output.  The kernel writes every element (live
  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());

-  VLLM_STABLE_DISPATCH_HALF_TYPES(
+  VLLM_DISPATCH_HALF_TYPES(
      q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
        using qkv_scalar_t = scalar_t;
        vllm::deepseek_v4_fused_ops::
            launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t const*>(q_in.const_data_ptr()),
-                reinterpret_cast<qkv_scalar_t*>(q_out.mutable_data_ptr()),
-                reinterpret_cast<qkv_scalar_t const*>(kv.const_data_ptr()),
-                reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-                slot_mapping.const_data_ptr<int64_t>(),
-                position_ids.const_data_ptr<int64_t>(),
-                cos_sin_cache.const_data_ptr<float>(), static_cast<float>(eps),
+                reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
+                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
+                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
+                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
+                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
                num_tokens_full, num_tokens_insert, num_heads_q,
                num_heads_q_padded, cache_block_size_i, kv_block_stride,
                stream);
      });
  return q_out;
 }
-
-// ────────────────────────────────────────────────────────────────────────────
-// FlashInfer full-cache torch ops
-// ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-    torch::stable::Tensor& q,                    // [N, H, 512] bf16, in place
-    torch::stable::Tensor const& kv,             // [N, 512] bf16, read-only
-    torch::stable::Tensor& k_cache,              // [num_blocks, bs, 512] bf16
-    torch::stable::Tensor const& slot_mapping,   // [num_tokens_insert] int64
-    torch::stable::Tensor const& position_ids,   // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, 64] float32
-    double eps, int64_t cache_block_size) {
-  using torch::headeronly::ScalarType;
-  STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
-                  "q must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() == ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() == ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
-                      cos_sin_cache.scalar_type() == ScalarType::Float &&
-                      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64] float32");
-  STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q.scalar_type() == ScalarType::BFloat16 &&
-                      kv.scalar_type() == ScalarType::BFloat16,
-                  "q and kv must be bfloat16");
-  STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
-                      k_cache.size(2) == 512 && k_cache.stride(2) == 1,
-                  "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
-  STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::BFloat16,
-                  "k_cache must be bfloat16");
-
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
-
-  // bf16 cache: 2 bytes/element -> byte strides for the uint8-addressed kernel.
-  int64_t const kv_block_stride = k_cache.stride(0) * 2;
-  int64_t const kv_token_stride = k_cache.stride(1) * 2;
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      q.scalar_type(),
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", [&] {
-        vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false,
-                                                           false>(
-            reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), nullptr, 0, 0,
-            reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
-            reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-            slot_mapping.const_data_ptr<int64_t>(),
-            position_ids.const_data_ptr<int64_t>(),
-            cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
-            static_cast<float>(eps), num_tokens_full, num_tokens_insert,
-            num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
-            kv_token_stride,
-            "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
-            stream);
-      });
-}
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-    torch::stable::Tensor const& q,                // [N, H, 512] bf16, read-only
-    torch::stable::Tensor const& kv,               // [N, 512] bf16, read-only
-    torch::stable::Tensor& q_fp8,                  // [N, H, 512] fp8 e4m3
-    torch::stable::Tensor& k_cache,                // [num_blocks, bs, 512] fp8
-    torch::stable::Tensor const& slot_mapping,     // [num_tokens_insert] int64
-    torch::stable::Tensor const& position_ids,     // [N] int64
-    torch::stable::Tensor const& cos_sin_cache,    // [max_pos, 64] float32
-    torch::stable::Tensor const& fp8_scale,        // scalar float32 (KV scale)
-    torch::stable::Tensor const& q_fp8_scale_inv,  // scalar float32 (1 / Q scale)
-    double eps, int64_t cache_block_size) {
-  using torch::headeronly::ScalarType;
-  STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
-                  "q must be contiguous CUDA");
-  STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
-                  "kv must be contiguous CUDA");
-  STD_TORCH_CHECK(q_fp8.device().is_cuda() && q_fp8.is_contiguous() &&
-                      q_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
-                      q_fp8.dim() == 3 && q_fp8.size(0) == q.size(0) &&
-                      q_fp8.size(1) == q.size(1) && q_fp8.size(2) == q.size(2),
-                  "q_fp8 must be a contiguous float8_e4m3fn tensor matching q");
-  STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
-  STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
-                      slot_mapping.scalar_type() == ScalarType::Long,
-                  "slot_mapping must be int64 CUDA");
-  STD_TORCH_CHECK(position_ids.device().is_cuda() &&
-                      position_ids.scalar_type() == ScalarType::Long,
-                  "position_ids must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
-                      cos_sin_cache.scalar_type() == ScalarType::Float &&
-                      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
-                  "cos_sin_cache shape [max_pos, 64] float32");
-  STD_TORCH_CHECK(fp8_scale.device().is_cuda() &&
-                      fp8_scale.scalar_type() == ScalarType::Float &&
-                      fp8_scale.size(0) == 1,
-                  "fp8_scale must be a scalar float32 CUDA tensor");
-  STD_TORCH_CHECK(q_fp8_scale_inv.device().is_cuda() &&
-                      q_fp8_scale_inv.scalar_type() == ScalarType::Float &&
-                      q_fp8_scale_inv.size(0) == 1,
-                  "q_fp8_scale_inv must be a scalar float32 CUDA tensor");
-  STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
-  STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  STD_TORCH_CHECK(q.scalar_type() == kv.scalar_type(),
-                  "q and kv dtype must match");
-  STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
-                      k_cache.size(2) == 512 && k_cache.stride(2) == 1,
-                  "k_cache shape [num_blocks, cache_block_size, 512] contiguous");
-  STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
-                  "k_cache must be float8_e4m3fn");
-
-  int const num_tokens_full = static_cast<int>(q.size(0));
-  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
-  STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
-                      static_cast<int>(position_ids.size(0)) == num_tokens_full,
-                  "q/kv/position_ids row counts must match");
-  STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
-                  "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      q.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      q.scalar_type(),
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", [&] {
-        vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, true,
-                                                           true>(
-            // q is read-only in the fp8 path (the kernel writes q_fp8); the
-            // launcher signature is non-const, so cast away const on the ptr.
-            reinterpret_cast<scalar_t*>(
-                const_cast<void*>(q.const_data_ptr())),
-            reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
-            q_fp8.stride(0), q_fp8.stride(1),
-            reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
-            reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
-            slot_mapping.const_data_ptr<int64_t>(),
-            position_ids.const_data_ptr<int64_t>(),
-            cos_sin_cache.const_data_ptr<float>(),
-            fp8_scale.const_data_ptr<float>(),
-            q_fp8_scale_inv.const_data_ptr<float>(), static_cast<float>(eps),
-            num_tokens_full, num_tokens_insert, num_heads_q,
-            static_cast<int>(cache_block_size),
-            // fp8 cache: 1 byte/element -> stride already in bytes.
-            k_cache.stride(0), k_cache.stride(1),
-            "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
-            stream);
-      });
-}
@@ -1,6 +1,9 @@
 #ifndef CONCAT_MLA_Q_CUH_
 #define CONCAT_MLA_Q_CUH_

+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
 #include "cuda_vec_utils.cuh"

 namespace vllm {
@@ -1,76 +0,0 @@
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-#include <torch/csrc/stable/c/shim.h>
-#include <torch/headeronly/version.h>
-#include <cuda_runtime.h>
-
-#include <array>
-#include <optional>
-
-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
-torch::stable::Tensor get_cuda_view_from_cpu_tensor(
-    torch::stable::Tensor& cpu_tensor) {
-  STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
-
-  const auto dtype = cpu_tensor.scalar_type();
-  const auto layout = cpu_tensor.layout();
-  const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA);
-
-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev);
-  }
-
-  std::array<StableIValue, 2> is_pinned_stack{
-      torch::stable::detail::from(cpu_tensor),
-      torch::stable::detail::from(std::nullopt)};
-  TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
-      "aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
-  if (torch::stable::detail::to<bool>(is_pinned_stack[0])) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ",
-                    cudaGetErrorString(err));
-
-    return torch::stable::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype,
-        [base = cpu_tensor](void*) {});  // keep cpu tensor alive
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor);
-  size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
-
-  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    STD_TORCH_CHECK(
-        false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
-
-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
-
-  return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(),
-                                  contiguous_cpu.strides(), cuda_dev,
-                                  contiguous_cpu.scalar_type(), deleter);
-}
@@ -1,69 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
-
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a,
-                              const torch::stable::Tensor& b,
-                              const torch::stable::Tensor& sfa,
-                              const torch::stable::Tensor& sfb,
-                              torch::stable::Tensor& d,
-                              const torch::stable::Tensor& problem_sizes,
-                              const torch::stable::Tensor& expert_offsets,
-                              const torch::stable::Tensor& blockscale_offsets) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-  STD_TORCH_CHECK(a.dim() == 2,
-                  "a must be a 2D tensor of shape (num_tokens, k)");
-  STD_TORCH_CHECK(b.dim() == 3,
-                  "b must be a 3D tensor of shape (num_experts, k, n)");
-  STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
-                  "k should align 128");
-  STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
-  STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major");
-  STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      a.get_device_index());
-  auto stream = get_current_cuda_stream(a.get_device_index());
-  if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                             blockscale_offsets, stream);
-  } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                         blockscale_offsets, stream);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented cutlass_mxfp8_grouped_mm for "
-                  "current device");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm));
-}
@@ -1,66 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "mxfp8_experts_quant.cuh"
-
-void mxfp8_experts_quant(const torch::stable::Tensor& input,
-                         const torch::stable::Tensor& problem_sizes,
-                         const torch::stable::Tensor& expert_offsets,
-                         const torch::stable::Tensor& blockscale_offsets,
-                         torch::stable::Tensor& quant_output,
-                         torch::stable::Tensor& scale_factor) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
-  STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
-  STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major");
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-
-  auto groups = problem_sizes.size(0);
-  STD_TORCH_CHECK(
-      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
-      "expert_offsets must be 1D and have size equal to the number of groups");
-  STD_TORCH_CHECK(
-      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
-      "blockscale_offsets must be 1D and have size equal to the number of "
-      "groups");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::launch_mxfp8_experts_quant<__half>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented mxfp8_experts_quant for "
-                  "current device");
-#endif
-}
-
-// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM
-// is applied only under COMPILE_LANGUAGE:CUDA.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant));
-}
@@ -3,6 +3,10 @@
 #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);
+
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                               torch::stable::Tensor& output_q,
                               torch::stable::Tensor& output_s,
@@ -24,10 +28,6 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
                                int64_t group_size, double eps, double int8_min,
                                double int8_max);

-torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
-                                   torch::stable::Tensor const& perm);
-
-#ifndef USE_ROCM
 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -162,11 +162,6 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 // AllSpark ops: declarations are in the source files
 // (allspark_repack.cu and allspark_qgemm_w8a16.cu)

-// TODO: Move this out once ROCm upgrade their torch to 2.11.
-// CPU tensor -> CUDA UVA view (shared CUDA)
-torch::stable::Tensor get_cuda_view_from_cpu_tensor(
-    torch::stable::Tensor& cpu_tensor);
-
 #endif

 // Attention kernels (shared CUDA/ROCm)
@@ -220,13 +215,6 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
                              std::optional<torch::stable::Tensor> residual,
                              int64_t group_size, bool is_scale_transposed);

-void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
-                                  torch::stable::Tensor const& input,
-                                  torch::stable::Tensor& scales,
-                                  int64_t group_size,
-                                  std::optional<torch::stable::Tensor> scale_ub,
-                                  bool is_scale_transposed);
-
 // Positional encoding kernels (shared CUDA/ROCm)
 void rotary_embedding(torch::stable::Tensor& positions,
                      torch::stable::Tensor& query,
@@ -243,44 +231,6 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                        torch::stable::Tensor& position_ids,
                        int64_t forced_token_heads_per_warp);

-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
-    torch::stable::Tensor& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, double eps,
-    int64_t cache_block_size);
-
-void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
-    torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
-    torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache,
-    torch::stable::Tensor const& fp8_scale,
-    torch::stable::Tensor const& q_fp8_scale_inv, double eps,
-    int64_t cache_block_size);
-
-#ifndef USE_ROCM
-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps);
-std::tuple<torch::stable::Tensor, torch::stable::Tensor>
-minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
-                         torch::stable::Tensor const& norm_weight_q,
-                         torch::stable::Tensor const& norm_weight_k,
-                         torch::stable::Tensor workspace, int64_t const q_size,
-                         int64_t const kv_size, int64_t const rank,
-                         int64_t const nranks, double const eps);
-#endif
-
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
    torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -323,26 +273,6 @@ void selective_scan_fwd(
    const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
    const std::optional<torch::stable::Tensor>& last_chunk_indices);

-using fptr_t = int64_t;
-fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
-                      torch::stable::Tensor& rank_data, int64_t rank,
-                      bool fully_connected);
-void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
-                torch::stable::Tensor& out, fptr_t reg_buffer,
-                int64_t reg_buffer_sz_bytes);
-void dispose(fptr_t _fa);
-int64_t meta_size();
-void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
-std::tuple<std::vector<int64_t>, std::vector<int64_t>>
-get_graph_buffer_ipc_meta(fptr_t _fa);
-void register_graph_buffers(fptr_t _fa,
-                            const std::vector<std::vector<int64_t>>& handles,
-                            const std::vector<std::vector<int64_t>>& offsets);
-std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
-    int64_t size);
-int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
-void free_shared_buffer(int64_t buffer);
-
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
@@ -18,7 +18,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "cutlass_extensions/torch_utils.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #include "get_group_starts.cuh"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
@@ -21,7 +21,7 @@
 #include "cutlass/util/packed_stride.hpp"
 #include "cutlass/util/mixed_dtype_utils.hpp"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 #include <cuda_runtime.h>
@@ -12,7 +12,7 @@

 #include <cutlass/arch/arch.h>

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #include "cute/tensor.hpp"
 #include "cutlass/tensor_ref.h"
@@ -20,7 +20,7 @@

 #include <cutlass/arch/arch.h>

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #include "cute/tensor.hpp"
 #include "cutlass/tensor_ref.h"
@@ -18,7 +18,7 @@

 #include "libtorch_stable/torch_utils.h"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "nvfp4_utils.cuh"

 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
@@ -18,7 +18,7 @@

 #include "libtorch_stable/torch_utils.h"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
 void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
@@ -18,7 +18,7 @@

 #include "libtorch_stable/torch_utils.h"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #include "cutlass/cutlass.h"

@@ -18,7 +18,7 @@

 #include "libtorch_stable/torch_utils.h"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 #include "cutlass/cutlass.h"

@@ -20,7 +20,7 @@
 #include "cutlass/util/packed_stride.hpp"

 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on

 namespace vllm::c3x {
@@ -15,7 +15,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"

 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on

 /*
@@ -1,7 +1,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/core/ScalarType.h>
 #include "cuda_utils.h"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
 void dispatch_scaled_mm(torch::stable::Tensor& c,
@@ -8,7 +8,7 @@

 #include <torch/csrc/stable/ops.h>
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 #include "get_group_starts.cuh"

 using namespace cute;
@@ -23,7 +23,7 @@
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"

 #include "core/math.hpp"
-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"
 // clang-format on

 using namespace cute;
@@ -4,7 +4,7 @@

 #include "libtorch_stable/torch_utils.h"

-#include "libtorch_stable/cutlass_extensions/common.hpp"
+#include "cutlass_extensions/common.hpp"

 void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
                            torch::stable::Tensor const& a,
@@ -7,11 +7,7 @@

 #include <cmath>

-#ifdef USE_ROCM
-  #include <hip/hip_fp8.h>
-#else
-  #include <cuda_fp8.h>
-#endif
+#include <cuda_fp8.h>

 #include "libtorch_stable/quantization/vectorization.cuh"
 #include "libtorch_stable/quantization/vectorization_utils.cuh"
@@ -19,23 +15,12 @@
 #include "libtorch_stable/torch_utils.h"

 __device__ __forceinline__ float GroupReduceMax(float val) {
-#ifdef USE_ROCM
-  // 16-thread logical groups may pack up to four per 64-lane wavefront; use a
-  // 64-bit mask and explicit width so shuffles stay within each group.
-  const int lane_in_wave = threadIdx.x % warpSize;
-  const unsigned long long mask = 0xFFFFull << ((lane_in_wave / 16) * 16);
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 8, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 16));
-  val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 16));
-#else
  unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;

  val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
  val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
  val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
  val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
-#endif
  return val;
 }

@@ -252,12 +237,10 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,

  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
      input.scalar_type(), "per_token_group_quant_8bit", ([&] {
-        if (dst_type == torch::headeronly::ScalarType::Char) {
+        if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
+          LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);
+        } else if (dst_type == torch::headeronly::ScalarType::Char) {
          LAUNCH_KERNEL(scalar_t, int8_t);
-        } else {
-          VLLM_STABLE_DISPATCH_FP8_TYPES(
-              dst_type, "per_token_group_quant_8bit_fp8",
-              ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));
        }
      }));

@@ -334,18 +317,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(

  // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
  // 8 lanes within the warp that share a group.
-#ifdef USE_ROCM
-  const int lane_in_wave = threadIdx.x % warpSize;
-  const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));
-  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));
-#else
  unsigned mask = 0xffu << (threadIdx.x & 24u);
  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
-#endif

  float y_s = local_absmax / max_8bit;
  y_s = fmaxf(y_s, 1e-10f);
@@ -528,12 +503,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,

  VLLM_STABLE_DISPATCH_HALF_TYPES(
      input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
-        if (dst_type == torch::headeronly::ScalarType::Char) {
+        if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
+          LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
+        } else if (dst_type == torch::headeronly::ScalarType::Char) {
          LAUNCH_REG_KERNEL(scalar_t, int8_t);
        } else {
-          VLLM_STABLE_DISPATCH_FP8_TYPES(
-              dst_type, "per_token_group_quant_8bit_packed_fp8",
-              ([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); }));
+          STD_TORCH_CHECK(
+              false,
+              "per_token_group_quant_8bit_packed only supports FP8/INT8 "
+              "outputs.");
        }
      }));

@@ -1,5 +1,4 @@
 #include "ops.h"
-#include "cuda_utils.h"
 #include "core/registration.h"

 #include <torch/csrc/stable/library.h>
@@ -8,6 +7,11 @@
 // 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, ops) {
+#ifndef USE_ROCM
+  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
+#endif
+
+#ifndef USE_ROCM
  // Compute per-token-group FP8 quantized tensor and scaling factor.
  // The dummy arguments are here so we can correctly fuse with RMSNorm.
  ops.def(
@@ -27,15 +31,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
      "output_s, int group_size, float eps, float int8_min, float int8_max) -> "
      "()");
-  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");

-#ifndef USE_ROCM
-
-  // TODO: Remove this once ROCm upgrade to torch 2.11.
-  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
-#endif
-
-#ifndef USE_ROCM
  // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
  // quantization, as well as bias
  ops.def(
@@ -325,16 +321,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "Tensor? scale_ub, Tensor!? residual, int group_size, "
      "bool is_scale_transposed) -> ()");

-  // Fused SiLU+Mul + per-block quantization
-  ops.def(
-      "silu_and_mul_per_block_quant("
-      "Tensor! out, "
-      "Tensor input, "
-      "Tensor! scales, "
-      "int group_size, "
-      "Tensor? scale_ub=None, "
-      "bool is_scale_transposed=False) -> ()");
-
  // Rotary embedding
  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
  ops.def(
@@ -351,38 +337,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "bool is_neox, Tensor position_ids, "
      "int forced_token_heads_per_warp=-1) -> ()");

-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor q_in, Tensor kv, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
-
-  // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
-  // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
-      "Tensor position_ids, Tensor cos_sin_cache, float eps, "
-      "int cache_block_size) -> ()");
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
-      "Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
-      "int cache_block_size) -> ()");
-
-#ifndef USE_ROCM
-  ops.def(
-      "minimax_allreduce_rms("
-      "Tensor input, Tensor norm_weight, Tensor workspace, "
-      "int rank, int nranks, float eps) -> Tensor");
-  ops.def(
-      "minimax_allreduce_rms_qk("
-      "Tensor qkv, Tensor norm_weight_q, Tensor norm_weight_k, "
-      "Tensor workspace, int q_size, int kv_size, int rank, int nranks, "
-      "float eps) -> (Tensor, Tensor)");
-#endif
-
  // Apply repetition penalties to logits in-place.
  ops.def(
      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -554,6 +508,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 }

 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
+#ifndef USE_ROCM
+  ops.impl("permute_cols", TORCH_BOX(&permute_cols));
+#endif
+
+#ifndef USE_ROCM
  // Per-token group quantization
  ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
  ops.impl("per_token_group_fp8_quant_packed",
@@ -561,9 +520,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("per_token_group_quant_int8",
           TORCH_BOX(&per_token_group_quant_int8));

-  ops.impl("permute_cols", TORCH_BOX(&permute_cols));
-
-#ifndef USE_ROCM
  // CUTLASS scaled_mm ops
  ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));
  ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp));
@@ -611,24 +567,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("rms_norm_dynamic_per_token_quant",
           TORCH_BOX(&rms_norm_dynamic_per_token_quant));
  ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
-  ops.impl("silu_and_mul_per_block_quant",
-           TORCH_BOX(&silu_and_mul_per_block_quant));

  // Positional encoding kernels (shared CUDA/ROCm)
  ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
  ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
-  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
-           TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
-  ops.impl(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
-      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
-  ops.impl(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
-      TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
-#ifndef USE_ROCM
-  ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
-  ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
-#endif

  // Sampler kernels (shared CUDA/ROCm)
  ops.impl("apply_repetition_penalties_",
@@ -675,28 +617,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }

-// TODO: Remove this once ROCm upgrade to torch 2.11.
-#ifndef USE_ROCM
-STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
-  ops.impl("get_cuda_view_from_cpu_tensor",
-           TORCH_BOX(&get_cuda_view_from_cpu_tensor));
-}
-
-STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
-  cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
-  cuda_utils.def(
-      "get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
-                          cuda_utils) {
-  cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
-  cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
-                  TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
-}
-
-#endif
-
 // These capability-check functions take only primitive args (no tensors), so
 // there is no device to dispatch on. CompositeExplicitAutograd makes them
 // available for all backends. This is the stable ABI equivalent of calling
@@ -805,45 +725,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
 }

-STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ar) {
-  custom_ar.def(
-      "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
-      "int rank, bool fully_connected) -> int");
-  custom_ar.def(
-      "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
-      "int reg_buffer_sz_bytes) -> ()");
-  custom_ar.def("dispose(int fa) -> ()");
-  custom_ar.def("meta_size() -> int");
-  custom_ar.def("register_buffer(int fa, int[] ipc_tensors) -> ()");
-  custom_ar.def("get_graph_buffer_ipc_meta(int fa) -> (int[], int[])");
-  custom_ar.def(
-      "register_graph_buffers(int fa, int[][] handles, int[][] offsets) -> ()");
-  custom_ar.def("allocate_shared_buffer_and_handle(int size) -> (int, Tensor)");
-  custom_ar.def("open_mem_handle(Tensor mem_handle) -> int");
-  custom_ar.def("free_shared_buffer(int ptr) -> ()");
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ar) {
-  custom_ar.impl("init_custom_ar", TORCH_BOX(&init_custom_ar));
-  custom_ar.impl("all_reduce", TORCH_BOX(&all_reduce));
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CPU, custom_ar) {
-  custom_ar.impl("open_mem_handle", TORCH_BOX(&open_mem_handle));
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CompositeExplicitAutograd, custom_ar) {
-  custom_ar.impl("dispose", TORCH_BOX(&dispose));
-  custom_ar.impl("meta_size", TORCH_BOX(&meta_size));
-  custom_ar.impl("register_buffer", TORCH_BOX(&register_buffer));
-  custom_ar.impl("get_graph_buffer_ipc_meta",
-                 TORCH_BOX(&get_graph_buffer_ipc_meta));
-  custom_ar.impl("register_graph_buffers", TORCH_BOX(&register_graph_buffers));
-  custom_ar.impl("allocate_shared_buffer_and_handle",
-                 TORCH_BOX(&allocate_shared_buffer_and_handle));
-  custom_ar.impl("free_shared_buffer", TORCH_BOX(&free_shared_buffer));
-}
-
 STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
  ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
 }
@@ -6,7 +6,11 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/shim_utils.h>

-#include <cuda_runtime.h>
+#ifndef USE_ROCM
+  #include <cuda_runtime.h>
+#else
+  #include <hip/hip_runtime.h>
+#endif
 #include <cublas_v2.h>

 #include <deque>
@@ -15,19 +15,16 @@
 * limitations under the License.
 */

-#include "torch_utils.h"
-
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/accelerator.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/csrc/stable/device.h>
-
 #include <cooperative_groups.h>
 #include <cuda_runtime.h>

+#include <torch/cuda.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
 #include "cuda_compat.h"
+#include "cuda_utils.h"
+#include "core/registration.h"
 #include "minimax_reduce_rms_kernel.h"

 #include <algorithm>
@@ -614,7 +611,7 @@ int get_sm_count() {
  static int sm_count = 0;
  if (sm_count == 0) {
    int device_id;
-    STD_CUDA_CHECK(cudaGetDevice(&device_id));
+    CUDA_CHECK(cudaGetDevice(&device_id));
    cudaDeviceProp device_prop;
    cudaGetDeviceProperties(&device_prop, device_id);
    sm_count = device_prop.multiProcessorCount;
@@ -624,13 +621,13 @@ int get_sm_count() {

 inline int getSMVersion(bool queryRealSmArch = false) {
  int device{-1};
-  STD_CUDA_CHECK(cudaGetDevice(&device));
+  CUDA_CHECK(cudaGetDevice(&device));
  int sm_major = 0;
  int sm_minor = 0;
-  STD_CUDA_CHECK(cudaDeviceGetAttribute(
-      &sm_major, cudaDevAttrComputeCapabilityMajor, device));
-  STD_CUDA_CHECK(cudaDeviceGetAttribute(
-      &sm_minor, cudaDevAttrComputeCapabilityMinor, device));
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_major,
+                                    cudaDevAttrComputeCapabilityMajor, device));
+  CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor,
+                                    cudaDevAttrComputeCapabilityMinor, device));
  int sm = sm_major * 10 + sm_minor;
  if (sm == 121 && !queryRealSmArch) {
    return 120;
@@ -642,7 +639,7 @@ template <typename KernelFunc>
 int get_max_active_blocks(KernelFunc kernel, int block_size,
                          int dynamic_smem = 0) {
  int max_active = 0;
-  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
      &max_active, kernel, block_size, dynamic_smem));
  return std::max(max_active, 1);
 }
@@ -681,27 +678,27 @@ void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
  cfg.attrs = attribute;
  cfg.numAttrs = SM >= 90 ? 2 : 0;

-  STD_CUDA_CHECK(cudaLaunchKernelEx(
+  CUDA_CHECK(cudaLaunchKernelEx(
      &cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
 }

 template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
 void minimax_reduce_rms_kernel_launcher_float4(
    MiniMaxReduceRMSParams const& params) {
-  STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0);
-  STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
+  TORCH_CHECK(params.size_q % params.hidden_dim == 0);
+  TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
  if (params.stride_q > 0) {
-    STD_TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
+    TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
  }
-  STD_TORCH_CHECK(params.allreduce_in_k != nullptr,
-                  "float4 QK kernel requires K input");
-  STD_TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
-  STD_TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
-  STD_TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
-  STD_TORCH_CHECK(params.size_q / params.hidden_dim ==
-                  params.size_k / params.hidden_dim_k);
+  TORCH_CHECK(params.allreduce_in_k != nullptr,
+              "float4 QK kernel requires K input");
+  TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
+  TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
+  TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
+  TORCH_CHECK(params.size_q / params.hidden_dim ==
+              params.size_k / params.hidden_dim_k);
  if (params.stride_k > 0) {
-    STD_TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
+    TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
  }

  int token_num = params.size_q / params.hidden_dim;
@@ -749,7 +746,7 @@ void minimax_reduce_rms_kernel_launcher_float4(
  cfg.attrs = attribute;
  cfg.numAttrs = SM >= 90 ? 2 : 0;

-  STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
+  CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
 }

 template <int NRanks>
@@ -762,21 +759,21 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
                    (params.hidden_dim * params.nranks == 6144) &&
                    (params.hidden_dim_k * params.nranks == 1024);

-  if (params.dtype == torch::headeronly::ScalarType::Half) {
+  if (params.dtype == at::ScalarType::Half) {
    if (use_float4) {
      minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
          params);
    } else {
      minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
    }
-  } else if (params.dtype == torch::headeronly::ScalarType::BFloat16) {
+  } else if (params.dtype == at::ScalarType::BFloat16) {
    if (use_float4) {
      minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
                                                1024>(params);
    } else {
      minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
    }
-  } else if (params.dtype == torch::headeronly::ScalarType::Float) {
+  } else if (params.dtype == at::ScalarType::Float) {
    if (use_float4) {
      minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
          params);
@@ -784,7 +781,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
      minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
    }
  } else {
-    STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
+    TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
  }
 }

@@ -798,18 +795,16 @@ void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
  } else if (params.nranks == 16) {
    dispatch_dtype<16>(params);
  } else {
-    STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
+    TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
  }
 }
 }  // namespace tensorrt_llm
 }  // namespace vllm

-torch::stable::Tensor minimax_allreduce_rms(
-    torch::stable::Tensor const& input,
-    torch::stable::Tensor const& norm_weight, torch::stable::Tensor workspace,
-    int64_t const rank, int64_t const nranks, double const eps) {
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps) {
  auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();

  allreduce_params.nranks = static_cast<int>(nranks);
@@ -820,12 +815,12 @@ torch::stable::Tensor minimax_allreduce_rms(
  allreduce_params.stride_q = allreduce_params.hidden_dim;
  allreduce_params.workspace =
      reinterpret_cast<void**>(workspace.mutable_data_ptr());
-  allreduce_params.allreduce_in = const_cast<void*>(input.const_data_ptr());
-  allreduce_params.rms_gamma = const_cast<void*>(norm_weight.const_data_ptr());
+  allreduce_params.allreduce_in = input.data_ptr();
+  allreduce_params.rms_gamma = norm_weight.data_ptr();
  allreduce_params.rms_eps = static_cast<float>(eps);
-  allreduce_params.stream = get_current_cuda_stream(input.get_device_index());
+  allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device());

-  torch::stable::Tensor rms_norm_out = torch::stable::empty_like(input);
+  torch::Tensor rms_norm_out = torch::empty_like(input);
  allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr();

  vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params);
@@ -833,33 +828,26 @@ torch::stable::Tensor minimax_allreduce_rms(
  return rms_norm_out;
 }

-std::tuple<torch::stable::Tensor, torch::stable::Tensor>
-minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
-                         torch::stable::Tensor const& norm_weight_q,
-                         torch::stable::Tensor const& norm_weight_k,
-                         torch::stable::Tensor workspace, int64_t const q_size,
-                         int64_t const kv_size, int64_t const rank,
-                         int64_t const nranks, double const eps) {
-  STD_TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
-  STD_TORCH_CHECK(qkv.is_contiguous(),
-                  "minimax_allreduce_rms_qk: qkv must be contiguous");
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps) {
+  TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
+  TORCH_CHECK(qkv.is_contiguous(),
+              "minimax_allreduce_rms_qk: qkv must be contiguous");
  int64_t qkv_dim = qkv.size(-1);
-  STD_TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
-                  "minimax_allreduce_rms_qk: qkv last dim must equal "
-                  "q_size + 2 * kv_size");
-  STD_TORCH_CHECK(rank < nranks,
-                  "minimax_allreduce_rms_qk: rank must be less than nranks");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      qkv.get_device_index());
+  TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
+              "minimax_allreduce_rms_qk: qkv last dim must equal "
+              "q_size + 2 * kv_size");
+  TORCH_CHECK(rank < nranks,
+              "minimax_allreduce_rms_qk: rank must be less than nranks");

  int64_t num_tokens = qkv.size(0);
  int elem_bytes = qkv.element_size();

-  torch::stable::Tensor q_out =
-      torch::stable::new_empty(qkv, {num_tokens, q_size}, qkv.scalar_type());
-  torch::stable::Tensor k_out =
-      torch::stable::new_empty(qkv, {num_tokens, kv_size}, qkv.scalar_type());
+  torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options());
+  torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options());

  auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
  params.nranks = static_cast<int>(nranks);
@@ -875,14 +863,13 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
  params.stride_k_out = 0;  // k_out is contiguous; kernel uses hidden_dim_k
  params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());

-  uint8_t* base =
-      const_cast<uint8_t*>(static_cast<const uint8_t*>(qkv.const_data_ptr()));
+  uint8_t* base = static_cast<uint8_t*>(qkv.data_ptr());
  params.allreduce_in = base;
  params.allreduce_in_k = base + q_size * elem_bytes;
-  params.rms_gamma = const_cast<void*>(norm_weight_q.const_data_ptr());
-  params.rms_gamma_k = const_cast<void*>(norm_weight_k.const_data_ptr());
+  params.rms_gamma = norm_weight_q.data_ptr();
+  params.rms_gamma_k = norm_weight_k.data_ptr();
  params.rms_eps = static_cast<float>(eps);
-  params.stream = get_current_cuda_stream(qkv.get_device_index());
+  params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device());

  params.rms_norm_out = q_out.mutable_data_ptr();
  params.rms_norm_out_k = k_out.mutable_data_ptr();
@@ -19,7 +19,7 @@
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>

-#include <torch/headeronly/core/ScalarType.h>
+#include <torch/types.h>

 namespace vllm {
 namespace tensorrt_llm {
@@ -51,7 +51,7 @@ static constexpr int kElemsPerAccess = ElemsPerAccess<DType>::value;
 struct MiniMaxReduceRMSParams {
  int nranks{};
  int rank{};
-  torch::headeronly::ScalarType dtype{torch::headeronly::ScalarType::Undefined};
+  at::ScalarType dtype{at::ScalarType::Undefined};
  int size_q{};
  int hidden_dim{};
  int size_k{};
@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
+
+#include <torch/all.h>
+
+#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
+
+void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
+                              const torch::Tensor& sfa,
+                              const torch::Tensor& sfb, torch::Tensor& d,
+                              const torch::Tensor& problem_sizes,
+                              const torch::Tensor& expert_offsets,
+                              const torch::Tensor& blockscale_offsets) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.size(1) == 3,
+              "problem_sizes must have shape (num_experts, 3)");
+  TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
+              "Number of experts in problem_sizes must match expert_offsets");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+  TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
+  TORCH_CHECK(b.dim() == 3,
+              "b must be a 3D tensor of shape (num_experts, k, n)");
+  TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
+              "k should align 128");
+  TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
+  TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
+  TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (d.dtype() == torch::kBFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                             blockscale_offsets, stream);
+  } else if (d.dtype() == torch::kFloat16) {
+    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
+        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
+                         blockscale_offsets, stream);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented cutlass_mxfp8_grouped_mm for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
+}
@@ -4,9 +4,9 @@
 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh

 #pragma once
-
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>

 #include <cassert>
 #include <iostream>
@@ -15,22 +15,18 @@
 #include "cute/tensor.hpp"
 #include "cutlass_mxfp8_grouped_mm_functor.cuh"
 #include "cutlass_mxfp8_grouped_mm_traits.cuh"
-#include "libtorch_stable/torch_utils.h"

 namespace expert_specialization {

 template <typename GemmTraits>
 void cutlass_mxfp8_grouped_mm_pre_compute(
-    torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs,
-    torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs,
-    torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a,
-    torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d,
-    torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
+    torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
+    torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
+    torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
+    torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
+    const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
  using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
  using ElementA = typename OffsetFunctor::ElementA;
  using ElementB = typename OffsetFunctor::ElementB;
@@ -46,10 +42,10 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
  using StrideB = typename StrideFunctor::StrideB;
  using StrideD = typename StrideFunctor::StrideD;

-  int num_experts = static_cast<int>(expert_offsets.size(0));
-  STD_TORCH_CHECK(num_experts <= 1024,
-                  "Number of experts cannot exceed 1024, the maximum number of "
-                  "threads per block.");
+  int num_experts = (int)expert_offsets.size(0);
+  TORCH_CHECK(num_experts <= 1024,
+              "Number of experts cannot exceed 1024, the maximum number of "
+              "threads per block.");

  OffsetFunctor offset_functor(
      reinterpret_cast<int*>(expert_offsets.data_ptr()),
@@ -76,18 +72,13 @@ void cutlass_mxfp8_grouped_mm_pre_compute(
 }

 template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
-                              const torch::stable::Tensor& b_ptrs,
-                              const torch::stable::Tensor& sfa_ptrs,
-                              const torch::stable::Tensor& sfb_ptrs,
-                              const torch::stable::Tensor& d_ptrs,
-                              const torch::stable::Tensor& stride_a,
-                              const torch::stable::Tensor& stride_b,
-                              const torch::stable::Tensor& stride_d,
-                              const torch::stable::Tensor& layout_sfa,
-                              const torch::stable::Tensor& layout_sfb,
-                              const torch::stable::Tensor& problem_sizes,
-                              cudaStream_t stream) {
+void cutlass_mxfp8_grouped_mm(
+    const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
+    const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
+    const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
+    const torch::Tensor& stride_b, const torch::Tensor& stride_d,
+    const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
+    const torch::Tensor& problem_sizes, cudaStream_t stream) {
  using Gemm = typename GemmTraits::Gemm;
  using ElementA = typename Gemm::ElementA;
  using ElementB = typename Gemm::ElementB;
@@ -102,12 +93,13 @@ void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
      typename GemmTraits::ProblemShape::UnderlyingProblemShape;

  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = d_ptrs.get_device_index();
-  hw_info.sm_count = get_device_prop()->multiProcessorCount;
+  hw_info.device_id = c10::cuda::current_device();
+  hw_info.sm_count =
+      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
  hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
  hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;

-  int num_experts = static_cast<int>(problem_sizes.size(0));
+  int num_experts = (int)problem_sizes.size(0);

  UnderlyingProblemShape* underlying_problem_shape =
      reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
@@ -135,55 +127,44 @@ void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
  Gemm gemm;

  auto can_implement_status = gemm.can_implement(arguments);
-  STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-                  "Failed to implement GEMM");
+  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
+              "Failed to implement GEMM");

+  torch::TensorOptions options_uint8 =
+      torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
  size_t workspace_size = gemm.get_workspace_size(arguments);
-  torch::stable::Tensor workspace = torch::stable::empty(
-      {static_cast<int64_t>(workspace_size)},
-      torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
+  torch::Tensor workspace = torch::empty(workspace_size, options_uint8);

  auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize GEMM");
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");

  status = gemm.run(stream, nullptr, true);  // Enable PDL
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }

 template <typename OutType>
 void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto device = a.device();
+    const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
+    const torch::Tensor& sfb, torch::Tensor& d,
+    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
+    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
+  int num_experts = (int)problem_sizes.size(0);
+  torch::TensorOptions options_int64 =
+      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
+  torch::TensorOptions options_int32 =
+      torch::TensorOptions().dtype(torch::kInt32).device(a.device());

-  torch::stable::Tensor a_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor b_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfa_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfb_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor d_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
+  torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
+  torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);

-  torch::stable::Tensor stride_a = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_b = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_d = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor layout_sfa =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
-  torch::stable::Tensor layout_sfb =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
+  torch::Tensor stride_a = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_b = torch::empty(num_experts, options_int64);
+  torch::Tensor stride_d = torch::empty(num_experts, options_int64);
+  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
+  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);

  using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
  cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
@@ -195,4 +176,4 @@ void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
      layout_sfa, layout_sfb, problem_sizes, stream);
 }

-}  // namespace expert_specialization
+}  // namespace expert_specialization
@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+// Adapted from SGLang:
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
+
+#include <torch/all.h>
+
+#include "mxfp8_experts_quant.cuh"
+
+void mxfp8_experts_quant(const torch::Tensor& input,
+                         const torch::Tensor& problem_sizes,
+                         const torch::Tensor& expert_offsets,
+                         const torch::Tensor& blockscale_offsets,
+                         torch::Tensor& quant_output,
+                         torch::Tensor& scale_factor) {
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
+  TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
+  TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
+  TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
+  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
+  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
+              "problem_sizes must be int32");
+  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
+              "expert_offsets must be int32");
+  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
+              "blockscale_offsets must be int32");
+
+  auto groups = problem_sizes.size(0);
+  TORCH_CHECK(
+      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
+      "expert_offsets must be 1D and have size equal to the number of groups");
+  TORCH_CHECK(
+      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
+      "blockscale_offsets must be 1D and have size equal to the number of "
+      "groups");
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  if (input.dtype() == torch::kBFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else if (input.dtype() == torch::kFloat16) {
+    expert_specialization::launch_mxfp8_experts_quant<__half>(
+        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
+        scale_factor);
+  } else {
+    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
+  }
+#else
+  TORCH_CHECK(false,
+              "No implemented mxfp8_experts_quant for "
+              "current device");
+#endif
+}
+
+#include "core/registration.h"
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
+}
@@ -4,19 +4,16 @@
 // https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh

 #pragma once
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <cuda.h>
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
-
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
+#include <torch/all.h>

 #include <cuda/ptx>

 #include "cute/tensor.hpp"
-#include "libtorch_stable/torch_utils.h"

 namespace expert_specialization {

@@ -359,12 +356,12 @@ __global__ void mxfp8_experts_quant_kernel(
 }

 template <typename T_IN>
-void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
-                                const torch::stable::Tensor& problem_sizes,
-                                const torch::stable::Tensor& expert_offsets,
-                                const torch::stable::Tensor& blockscale_offsets,
-                                torch::stable::Tensor& quant_output,
-                                torch::stable::Tensor& scale_factor) {
+void launch_mxfp8_experts_quant(const torch::Tensor& input,
+                                const torch::Tensor& problem_sizes,
+                                const torch::Tensor& expert_offsets,
+                                const torch::Tensor& blockscale_offsets,
+                                torch::Tensor& quant_output,
+                                torch::Tensor& scale_factor) {
  ThrLayout thr_layout{};
  ValLayout val_layout{};
  SfR2SThrLayout r2s_thr_layout{};
@@ -389,18 +386,19 @@ void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
      CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)

  int max_active_blocks_per_sm = -1;
-  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+  AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
      &max_active_blocks_per_sm,
      mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
                                 decltype(tiled_copy_r2g),
                                 decltype(tiled_copy_r2s)>,
      THREAD_BLOCK_SIZE, 0));

-  dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
+  dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
+                max_active_blocks_per_sm,
            1, 1);
  dim3 block(THREAD_BLOCK_SIZE, 1, 1);
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto stream = get_current_cuda_stream(input.get_device_index());
+  int num_experts = (int)problem_sizes.size(0);
+  auto stream = at::cuda::getCurrentCUDAStream();
  mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
                             decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
      <<<grid, block, 0, stream>>>(
@@ -40,6 +40,18 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                        torch::Tensor& weight, double epsilon);

+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
+    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
+
+void silu_and_mul_per_block_quant(torch::Tensor& out,
+                                  torch::Tensor const& input,
+                                  torch::Tensor& scales, int64_t group_size,
+                                  std::optional<torch::Tensor> scale_ub,
+                                  bool is_scale_transposed);
+
 // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
 // ABI for CUDA). It remains here because the CPU build still uses these
 // torch::Tensor declarations.
@@ -78,6 +90,8 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
                        torch::Tensor const& seq_lens,
                        torch::Tensor const& page_table, double scale);

+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
+
 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor const& scale,
                              std::optional<torch::Tensor> const& azp);
@@ -93,6 +107,24 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
    int64_t activation_kind);

 using fptr_t = int64_t;
+fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
+                      torch::Tensor& rank_data, int64_t rank,
+                      bool fully_connected);
+void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
+                fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
+void dispose(fptr_t _fa);
+int64_t meta_size();
+void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa);
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets);
+std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
+    int64_t size);
+int64_t open_mem_handle(torch::Tensor& mem_handle);
+void free_shared_buffer(int64_t buffer);
+
 #ifdef USE_ROCM
 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
                      std::optional<int64_t> qr_max_size = std::nullopt);
@@ -102,7 +134,16 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
 void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                   int64_t quant_level, bool cast_bf2half = false);
 int64_t qr_max_size();
+#endif

-// TODO: Remove this once ROCm upgrade to torch 2.11.
-torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
+#ifndef USE_ROCM
+torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
+                                    torch::Tensor const& norm_weight,
+                                    torch::Tensor workspace, int64_t const rank,
+                                    int64_t const nranks, double const eps);
+std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
+    torch::Tensor qkv, torch::Tensor const& norm_weight_q,
+    torch::Tensor const& norm_weight_k, torch::Tensor workspace,
+    int64_t const q_size, int64_t const kv_size, int64_t const rank,
+    int64_t const nranks, double const eps);
 #endif
@@ -1,10 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-#include "../../torch_utils.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>

 #include "../../dispatch_utils.h"
-#include "quant_conversions.cuh"
+#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"

 namespace vllm {

@@ -104,70 +105,64 @@ __global__ void silu_and_mul_per_block_quant_kernel(

 }  // namespace vllm

-void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
-                                  torch::stable::Tensor const& input,
-                                  torch::stable::Tensor& scales,
-                                  int64_t group_size,
-                                  std::optional<torch::stable::Tensor> scale_ub,
+void silu_and_mul_per_block_quant(torch::Tensor& out,
+                                  torch::Tensor const& input,
+                                  torch::Tensor& scales, int64_t group_size,
+                                  std::optional<torch::Tensor> scale_ub,
                                  bool is_scale_transposed) {
-  static torch::headeronly::ScalarType kFp8Type =
-      is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
-                   : torch::headeronly::ScalarType::Float8_e4m3fnuz;
+  static c10::ScalarType kFp8Type = is_fp8_ocp()
+                                        ? c10::ScalarType::Float8_e4m3fn
+                                        : c10::ScalarType::Float8_e4m3fnuz;

-  STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
-                  out.scalar_type() == torch::headeronly::ScalarType::Char);
-  STD_TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
-  STD_TORCH_CHECK(
-      input.scalar_type() == torch::headeronly::ScalarType::Half ||
-          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+  TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
+  TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
+  TORCH_CHECK(
+      input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
      "Input must be FP16 or BF16");
-  STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
-  STD_TORCH_CHECK(group_size == 128 || group_size == 64,
-                  "Unsupported group size: ", group_size);
+  TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
+  TORCH_CHECK(group_size == 128 || group_size == 64,
+              "Unsupported group size: ", group_size);

  if (scale_ub.has_value()) {
-    STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
+    TORCH_CHECK(out.dtype() == kFp8Type);
  }

  int32_t hidden_size = out.size(-1);
  auto num_tokens = input.size(0);
  int32_t num_groups = hidden_size / group_size;

-  STD_TORCH_CHECK(input.size(-1) == hidden_size * 2,
-                  "input last dim must be 2x output hidden_size");
-  STD_TORCH_CHECK(hidden_size % group_size == 0,
-                  "hidden_size must be divisible by group_size");
+  TORCH_CHECK(input.size(-1) == hidden_size * 2,
+              "input last dim must be 2x output hidden_size");
+  TORCH_CHECK(hidden_size % group_size == 0,
+              "hidden_size must be divisible by group_size");

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  dim3 grid(num_tokens, num_groups);
  dim3 block(group_size);

-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
+  VLLM_DISPATCH_FLOATING_TYPES(
      input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
        using scalar_in_t = scalar_t;

-        VLLM_STABLE_DISPATCH_QUANT_TYPES(
+        VLLM_DISPATCH_QUANT_TYPES(
            out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
              using scalar_out_t = scalar_t;

-              VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
-                VLLM_STABLE_DISPATCH_BOOL(
-                    is_scale_transposed, transpose_scale, [&] {
-                      vllm::silu_and_mul_per_block_quant_kernel<
-                          scalar_in_t, scalar_out_t, transpose_scale, gs>
-                          <<<grid, block, 0, stream>>>(
-                              out.mutable_data_ptr<scalar_out_t>(),
-                              scales.mutable_data_ptr<float>(),
-                              input.const_data_ptr<scalar_in_t>(),
-                              scale_ub.has_value()
-                                  ? scale_ub->const_data_ptr<float>()
-                                  : nullptr,
-                              hidden_size);
-                    });
+              VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
+                VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
+                  vllm::silu_and_mul_per_block_quant_kernel<
+                      scalar_in_t, scalar_out_t, transpose_scale, gs>
+                      <<<grid, block, 0, stream>>>(
+                          out.data_ptr<scalar_out_t>(),
+                          scales.data_ptr<float>(),
+                          input.data_ptr<scalar_in_t>(),
+                          scale_ub.has_value() ? scale_ub->data_ptr<float>()
+                                               : nullptr,
+                          hidden_size);
+                });
              });
            });
      });
-}
+}
@@ -1,639 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100).
-//
-// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3
-// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch.
-// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering
-// a tile of N output columns and K input positions.
-//
-// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled,
-// [E, groups, N] scales, [E, groups, N/8] packed zeros.
-//
-// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread
-// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to
-// the pre-zeroed output tensor (no FP32 scratch buffer).
-
-#include <cstdint>
-
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
-
-#include "qdq_4_rdna3.cuh"
-
-#if defined(__HIPCC__) && defined(__gfx1100__)
-  #define __HIP__RDNA3__
-#endif
-
-namespace vllm {
-namespace moe_gptq_rdna3 {
-
-#define BLOCK_KN_SIZE 256
-#define THREADS_X 256
-
-#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
-
-using gptq_rdna3::bf162_t;
-using gptq_rdna3::bf16_t;
-
-// --- Helpers (same as q_gemm_rdna3.cu) ---
-
-template <typename T>
-__forceinline__ __device__ T tzero();
-
-template <>
-__forceinline__ __device__ half tzero<half>() {
-  return __float2half_rn(0.0f);
-}
-
-template <>
-__forceinline__ __device__ bf16_t tzero<bf16_t>() {
-  return __float2bfloat16(0.0f);
-}
-
-__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
-  float result = 0.0f;
-  const half2* a2_ptr = (const half2*)a_ptr;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
-  }
-  return result;
-}
-
-__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
-                                           const bf16_t* a_ptr) {
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
-    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
-  }
-  return result;
-}
-
-__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
-                                                   half2 v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      half2 h2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.h2[0] = __hadd2(cur.h2[0], v01);
-    sum.h2[1] = __hadd2(cur.h2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
-                                                    bf162_t v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      bf162_t b2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.b2[0] = __hadd2(cur.b2[0], v01);
-    sum.b2[1] = __hadd2(cur.b2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
-                                            int (&zeros)[4]) {
-  int qcol = n / 8;
-  int shift = (n & 0x07) * 4;
-  uint32_t d = qzeros_row[qcol] >> shift;
-  zeros[0] = (int)(d & 0xF);
-  zeros[1] = (int)((d >> 4) & 0xF);
-  zeros[2] = (int)((d >> 8) & 0xF);
-  zeros[3] = (int)((d >> 12) & 0xF);
-}
-
-template <typename T>
-__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
-                                             T (&scales)[4]) {
-  scales[0] = scales_row[n + 0];
-  scales[1] = scales_row[n + 1];
-  scales[2] = scales_row[n + 2];
-  scales[3] = scales_row[n + 3];
-}
-
-// ---------------------------------------------------------------------------
-// Fused MoE kernel.
-// ---------------------------------------------------------------------------
-
-template <typename T, int BLOCK_SIZE_M>
-__global__ void moe_gemm_q4_kernel_rdna3(
-    const T* __restrict__ a,                  // [size_m, size_k] or [M*topk, K]
-    T* __restrict__ c,                        // [M*topk, size_n] pre-zeroed
-    const uint32_t* __restrict__ b_q_weight,  // [E, K/8, N] packed
-    const T* __restrict__ b_scales,           // [E, groups, N]
-    const uint32_t* __restrict__ b_qzeros,    // [E, groups, N/8] packed
-    const float* __restrict__ topk_weights,   // [M*topk] or nullptr
-    const int32_t* __restrict__ sorted_token_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ num_tokens_post_padded,
-    const int size_m,  // total tokens (original M, or M*topk for w2)
-    const int size_n,  // output features per expert
-    const int size_k,  // input features
-    const int groups,  // K / group_size
-    const int top_k,   // routing top-k (1 for w2 pass)
-    // Per-expert strides (in elements, not bytes)
-    const int expert_weight_stride,  // (K/8) * N
-    const int expert_scales_stride,  // groups * N
-    const int expert_zeros_stride,   // groups * (N/8)
-    const bool mul_topk_weight,
-    const int output_topk) {  // >0: reduce output by token_id/output_topk
-  const int t = threadIdx.x;
-  const int token_block = blockIdx.x;
-  const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
-  const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
-  const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
-  const int n = offset_n + t * 4;
-
-  // Early exit for padding blocks or invalid experts (expert_map = -1)
-  if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
-
-  const int expert_id = expert_ids[token_block];
-  if (expert_id == -1) return;
-
-  // Expert-specific pointers
-  const uint32_t* expert_weights =
-      b_q_weight + (int64_t)expert_id * expert_weight_stride;
-  const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
-  const uint32_t* expert_qzeros =
-      b_qzeros + (int64_t)expert_id * expert_zeros_stride;
-
-  // LDS for activations
-  constexpr int LDS_PAD = 8;
-  __shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
-
-  static_assert(BLOCK_KN_SIZE == THREADS_X,
-                "BLOCK_KN_SIZE must equal THREADS_X");
-
-  // For bf16 M=1, we can skip LDS and read A from global (same as dense).
-  // fp16 always needs LDS due to the dot22_8_f indexing pattern.
-  constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
-
-  const int offset_m_base = token_block * BLOCK_SIZE_M;
-
-  if constexpr (USE_LDS_A) {
-    if (offset_k + t < end_k) {
-  #pragma unroll
-      for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-        int32_t token_id = sorted_token_ids[offset_m_base + m];
-        int token_row = token_id / top_k;
-        T av;
-        if (token_row < size_m) {
-          av = a[(int64_t)token_row * size_k + offset_k + t];
-        } else {
-          av = tzero<T>();
-        }
-        block_a[m][t] = av;
-      }
-    }
-    __syncthreads();
-  }
-
-  if (n >= size_n) return;
-
-  // Group bookkeeping
-  const int groupsize = size_k / groups;
-  int group = offset_k / groupsize;
-  int nextgroup = (group + 1) * groupsize;
-
-  // Weight pointer for this expert
-  int qk = offset_k / 8;
-  const uint32_t* b_ptr = expert_weights + qk * size_n + n;
-
-  // Per-column dequant constants (4 columns per thread)
-  half2 z1z16_h[4][2], y1y16_h[4][2];
-  float z_b_f[4], y_b_f[4];
-
-  // GPTQv1: zero_offset = 1
-  constexpr int zero_offset = 1;
-
-  auto refresh_group = [&](int g) {
-    const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
-    const T* sc_row = expert_scales + g * size_n;
-    int zeros[4];
-    T scales[4];
-    load4_zeros(qz_row, n, zeros);
-    load4_scales<T>(sc_row, n, scales);
-    if constexpr (std::is_same<T, half>::value) {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
-                                         scales[i], z1z16_h[i], y1y16_h[i]);
-      }
-    } else {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
-                                             scales[i], z_b_f[i], y_b_f[i]);
-      }
-    }
-  };
-
-  refresh_group(group);
-
-  float block_c[BLOCK_SIZE_M][4];
-  #pragma unroll
-  for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
-  }
-
-  // --- Main K-loop ---
-  int k = offset_k;
-  while (k < end_k) {
-    if (k == nextgroup) {
-      group++;
-      nextgroup += groupsize;
-      refresh_group(group);
-    }
-
-    // Prefetch 4 weight words (128 bytes)
-    int4 b_w[4];
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      b_w[j] = *(const int4*)(b_ptr + j * size_n);
-    }
-    b_ptr += 4 * size_n;
-
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      const int a_off = (k - offset_k) + 8 * j;
-
-      if constexpr (std::is_same<T, half>::value) {
-        // fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
-        half2 dq[4][4];
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
-                                        y1y16_h[0]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
-                                        y1y16_h[1]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
-                                        y1y16_h[2]);
-        gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
-                                        y1y16_h[3]);
-
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
-          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
-          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
-          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
-          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
-        }
-      } else if constexpr (BLOCK_SIZE_M == 1) {
-        // bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        // Load activations — read from global (no LDS for bf16 M=1)
-        pack4 a_pack;
-        {
-          int32_t token_id = sorted_token_ids[offset_m_base];
-          int token_row = token_id / top_k;
-          if (token_row < size_m) {
-            const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
-                a + (int64_t)token_row * size_k + offset_k + a_off);
-            a_pack.u[0] = a_words[0];
-            a_pack.u[1] = a_words[1];
-            a_pack.u[2] = a_words[2];
-            a_pack.u[3] = a_words[3];
-          } else {
-            a_pack.u[0] = 0;
-            a_pack.u[1] = 0;
-            a_pack.u[2] = 0;
-            a_pack.u[3] = 0;
-          }
-        }
-
-        // sum_a for bias correction
-        float sum_a = 0.0f;
-  #pragma unroll
-        for (int b = 0; b < 4; ++b) {
-          sum_a = __builtin_amdgcn_fdot2_f32_bf16(
-              *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
-              sum_a, /*clamp=*/false);
-        }
-
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-          float partial = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            partial = __builtin_amdgcn_fdot2_f32_bf16(
-                *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                partial, /*clamp=*/false);
-          }
-
-          block_c[0][col] =
-              __fmaf_rn(y_b_f[col], partial,
-                        __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
-        }
-      } else {
-        // bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        pack4 a_pack[BLOCK_SIZE_M];
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          const uint32_t* a_words =
-              reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
-          a_pack[m].u[0] = a_words[0];
-          a_pack[m].u[1] = a_words[1];
-          a_pack[m].u[2] = a_words[2];
-          a_pack[m].u[3] = a_words[3];
-        }
-
-        float sum_a[BLOCK_SIZE_M];
-  #pragma unroll
-        for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-          float s = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
-                                                *((const bf16x2_t*)&BF16_ONES),
-                                                s, /*clamp=*/false);
-          }
-          sum_a[m] = s;
-        }
-
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-  #pragma unroll
-          for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-            float partial = 0.0f;
-  #pragma unroll
-            for (int b = 0; b < 4; ++b) {
-              partial = __builtin_amdgcn_fdot2_f32_bf16(
-                  *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                  partial, /*clamp=*/false);
-            }
-            block_c[m][col] =
-                __fmaf_rn(y_b_f[col], partial,
-                          __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
-          }
-        }
-      }
-    }
-    k += 32;
-  }
-
-  // --- Epilogue: apply topk_weight and atomic-add to output ---
-  #pragma unroll
-  for (int m = 0; m < BLOCK_SIZE_M; ++m) {
-    int32_t token_id = sorted_token_ids[offset_m_base + m];
-    if (token_id / top_k >= size_m) continue;
-
-    // Apply router weight
-    if (mul_topk_weight && topk_weights != nullptr) {
-      float tw = topk_weights[token_id];
-  #pragma unroll
-      for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
-    }
-
-    // output_topk > 0: reduce by mapping token_id back to original token
-    // (multiple experts write to the same row via atomics)
-    int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
-                                        : (int64_t)token_id;
-    T* out = c + out_row * size_n + n;
-    if constexpr (std::is_same<T, half>::value) {
-      half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
-                                 __float2half_rn(block_c[m][1]));
-      half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
-                                 __float2half_rn(block_c[m][3]));
-      atomic_add_pk4_f16(out, r01, r23);
-    } else {
-      bf162_t r01;
-      r01.x = __float2bfloat16(block_c[m][0]);
-      r01.y = __float2bfloat16(block_c[m][1]);
-      bf162_t r23;
-      r23.x = __float2bfloat16(block_c[m][2]);
-      r23.y = __float2bfloat16(block_c[m][3]);
-      atomic_add_pk4_bf16(out, r01, r23);
-    }
-  }
-}
-
-#else  // non-RDNA3: empty stub for symbol parity
-
-template <typename T, int BLOCK_SIZE_M>
-__global__ void moe_gemm_q4_kernel_rdna3(
-    const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
-    const int32_t*, const int32_t*, const int32_t*, const int, const int,
-    const int, const int, const int, const int, const int, const int,
-    const bool, const int) {}
-
-#endif  // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
-
-// ---------------------------------------------------------------------------
-// Launcher
-// ---------------------------------------------------------------------------
-
-template <typename T, int BLOCK_SIZE_M>
-void launch_moe_gemm_q4(
-    const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
-    const uint32_t* b_qzeros, const float* topk_weights,
-    const int32_t* sorted_token_ids, const int32_t* expert_ids,
-    const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
-    int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
-    int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
-    int output_topk, cudaStream_t stream) {
-  dim3 block(THREADS_X);
-  dim3 grid(num_token_blocks,
-            (size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
-            (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
-
-  moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
-      a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-      expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
-      expert_weight_stride, expert_scales_stride, expert_zeros_stride,
-      mul_topk_weight, output_topk);
-}
-
-template <typename T>
-void dispatch_moe_gemm_q4(
-    const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
-    const uint32_t* b_qzeros, const float* topk_weights,
-    const int32_t* sorted_token_ids, const int32_t* expert_ids,
-    const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
-    int size_n, int size_k, int groups, int top_k, int block_size_m,
-    int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
-    bool mul_topk_weight, int output_topk, cudaStream_t stream) {
-  // Dispatch to template instantiation based on block_size_m
-  switch (block_size_m) {
-    case 1:
-      launch_moe_gemm_q4<T, 1>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 2:
-      launch_moe_gemm_q4<T, 2>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 4:
-      launch_moe_gemm_q4<T, 4>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    case 8:
-      launch_moe_gemm_q4<T, 8>(
-          a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
-          expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
-          size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
-          expert_zeros_stride, mul_topk_weight, output_topk, stream);
-      break;
-    default:
-      TORCH_CHECK(false,
-                  "moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
-                  "got ",
-                  block_size_m);
-  }
-}
-
-}  // namespace moe_gptq_rdna3
-}  // namespace vllm
-
-// ---------------------------------------------------------------------------
-// Public entry point
-// ---------------------------------------------------------------------------
-//
-// Inputs:
-//   a                      [M, K] or [M*top_k, K]  half or bfloat16
-//   c                      [M*top_k, N]             same dtype (pre-zeroed!)
-//   b_q_weight             [E, K/8, N]              uint32 (shuffled)
-//   b_scales               [E, groups, N]           same dtype as a
-//   b_qzeros               [E, groups, N/8]         uint32 (packed 4-bit)
-//   topk_weights           [M*top_k] or empty       float32
-//   sorted_token_ids       [num_blocks * block_m]   int32
-//   expert_ids             [num_blocks]              int32
-//   num_tokens_post_padded [1]                       int32
-//   top_k                  int
-//   block_size_m           int (1, 2, 4, or 8)
-//   mul_topk_weight        bool
-
-void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
-                         torch::Tensor b_q_weight, torch::Tensor b_scales,
-                         torch::Tensor b_qzeros, torch::Tensor topk_weights,
-                         torch::Tensor sorted_token_ids,
-                         torch::Tensor expert_ids,
-                         torch::Tensor num_tokens_post_padded, int64_t top_k,
-                         int64_t block_size_m, bool mul_topk_weight,
-                         int64_t output_topk) {
-  TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
-  TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
-  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
-  TORCH_CHECK(a.dim() == 2, "a must be 2D");
-  TORCH_CHECK(c.dim() == 2, "c must be 2D");
-  TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
-  TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
-  TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
-  TORCH_CHECK(
-      a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
-      "a must be half or bfloat16");
-  TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
-              "b_scales dtype must match a");
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  int size_m = (int)a.size(0);
-  int size_k = (int)a.size(1);
-  int size_n = (int)b_q_weight.size(2);
-  int groups = (int)b_scales.size(1);
-
-  // Per-expert strides
-  int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
-  int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
-  int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
-
-  int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
-
-  const float* topk_w_ptr =
-      (topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
-
-  // Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
-  // implicitly convert to half/__hip_bfloat16 in device code).
-  using vllm::gptq_rdna3::bf16_t;
-
-  auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
-    using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
-    vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
-        a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
-        (const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
-        sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
-        num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
-        size_n, size_k, groups, (int)top_k, (int)block_size_m,
-        expert_weight_stride, expert_scales_stride, expert_zeros_stride,
-        mul_topk_weight, (int)output_topk, stream);
-  };
-
-  if (a.scalar_type() == torch::kHalf) {
-    dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
-             (const half*)b_scales.data_ptr());
-  } else {
-    dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
-             (const bf16_t*)b_scales.data_ptr());
-  }
-}
@@ -27,15 +27,6 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
                                   torch::Tensor b_scales,
                                   torch::Tensor b_g_idx, bool use_v2_format);

-void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
-                         torch::Tensor b_q_weight, torch::Tensor b_scales,
-                         torch::Tensor b_qzeros, torch::Tensor topk_weights,
-                         torch::Tensor sorted_token_ids,
-                         torch::Tensor expert_ids,
-                         torch::Tensor num_tokens_post_padded, int64_t top_k,
-                         int64_t block_size_m, bool mul_topk_weight,
-                         int64_t output_topk);
-
 void paged_attention(
    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -50,15 +50,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
      "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
      "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
  rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
-
-  rocm_ops.def(
-      "moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
-      "Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
-      "Tensor sorted_token_ids, Tensor expert_ids, "
-      "Tensor num_tokens_post_padded, "
-      "int top_k, int block_size_m, bool mul_topk_weight, "
-      "int output_topk) -> ()");
-  rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
 #endif

  // Custom attention op
@@ -32,24 +32,37 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);

-#ifdef USE_ROCM
-  // TODO: Remove this once we upgrade to torch 2.11.
-  // ROCm still uses torch 2.10,
-  // So we still need to use unstable torch ABI for now.
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
  ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
           &get_cuda_view_from_cpu_tensor);
-#endif

  // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
  ops.def(
      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);

+  // Fused SiLU+Mul + per-block quantization
+  ops.def(
+      "silu_and_mul_per_block_quant("
+      "Tensor! out, "
+      "Tensor input, "
+      "Tensor! scales, "
+      "int group_size, "
+      "Tensor? scale_ub=None, "
+      "bool is_scale_transposed=False) -> ()");
+  ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
+           &silu_and_mul_per_block_quant);
+
  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
-  // full-cache bf16/fp8 variants).
+  // kernel launch.
+  ops.def(
+      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
+      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
+  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
+           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);

  // Quantization ops
 #ifndef USE_ROCM
@@ -150,29 +163,37 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // conditionally compiled so impl registration is in source file

 #endif
+
+#ifndef USE_ROCM
+  ops.def(
+      "minimax_allreduce_rms("
+      "Tensor input,"
+      "Tensor norm_weight,"
+      "Tensor workspace,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> Tensor");
+  ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms);
+  ops.def(
+      "minimax_allreduce_rms_qk("
+      "Tensor qkv,"
+      "Tensor norm_weight_q,"
+      "Tensor norm_weight_k,"
+      "Tensor workspace,"
+      "int q_size,"
+      "int kv_size,"
+      "int rank,"
+      "int nranks,"
+      "float eps) -> (Tensor, Tensor)");
+  ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
+
+  //  conditionally compiled so impl in source file
+#endif
 }

-#ifdef USE_ROCM
-TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
-  // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C).
-  custom_ar.def(
-      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
-      "cast_bf2half) -> ()");
-  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
-
-  custom_ar.def("init_custom_qr", &init_custom_qr);
-  custom_ar.def("qr_destroy", &qr_destroy);
-  custom_ar.def("qr_get_handle", &qr_get_handle);
-
-  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
-  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
-
-  custom_ar.def("qr_max_size", &qr_max_size);
-}
-
-// TODO: Remove this once ROCm upgrade to torch 2.11.
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
  // Cuda utils
+
  // Gets the specified device attribute.
  cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
  cuda_utils.impl("get_device_attribute", &get_device_attribute);
@@ -183,6 +204,49 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
  cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
                  &get_max_shared_memory_per_block_device_attribute);
 }
+
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
+  // Custom all-reduce kernels
+  custom_ar.def(
+      "init_custom_ar(int[] ipc_tensors, Tensor rank_data, "
+      "int rank, bool fully_connected) -> int");
+  custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar);
+  custom_ar.def(
+      "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
+      "int reg_buffer_sz_bytes) -> ()");
+  custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce);
+
+  custom_ar.def("dispose", &dispose);
+  custom_ar.def("meta_size", &meta_size);
+
+  custom_ar.def("register_buffer", &register_buffer);
+  custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta);
+  custom_ar.def("register_graph_buffers", &register_graph_buffers);
+
+  custom_ar.def("allocate_shared_buffer_and_handle",
+                &allocate_shared_buffer_and_handle);
+  custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle);
+  custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
+
+  custom_ar.def("free_shared_buffer", &free_shared_buffer);
+#ifdef USE_ROCM
+  // Quick Reduce all-reduce kernels
+  custom_ar.def(
+      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
+      "cast_bf2half) -> ()");
+  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
+
+  custom_ar.def("init_custom_qr", &init_custom_qr);
+  custom_ar.def("qr_destroy", &qr_destroy);
+
+  custom_ar.def("qr_get_handle", &qr_get_handle);
+
+  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
+  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
+
+  // Max input size in bytes
+  custom_ar.def("qr_max_size", &qr_max_size);
 #endif
+}

 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
@@ -50,7 +50,7 @@ struct _typeConvert<float> {
 #if defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000))
 // CUDA < 12.0 runs into issues with packed type conversion
 template <>
-struct _typeConvert<torch::headeronly::Half> {
+struct _typeConvert<c10::Half> {
  static constexpr bool exists = true;
  using hip_type = __half;
  using packed_hip_type = __half2;
@@ -73,7 +73,7 @@ struct _typeConvert<torch::headeronly::Half> {
 // CUDA_ARCH < 800 does not have BF16 support
 // ROCm 7.0+ supports bfloat16
 template <>
-struct _typeConvert<torch::headeronly::BFloat16> {
+struct _typeConvert<c10::BFloat16> {
  static constexpr bool exists = true;
  using hip_type = __nv_bfloat16;
  using packed_hip_type = __nv_bfloat162;
@@ -255,49 +255,51 @@ ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 #################### RUST BUILD IMAGE ####################
 # Build the Rust frontend (`vllm-rs`) in a dedicated stage so the main wheel
 # build stage doesn't need the rust toolchain, protoc, or the rust source.
-# This stage reuses the Python environment from base and runs in parallel with
-# csrc-build/extensions-build.
-FROM base AS rust-build
+# This stage runs in parallel with csrc-build/extensions-build.
+FROM ${BUILD_BASE_IMAGE} AS rust-build
 ARG BUILD_OS

-# Install native tools needed only for Rust/protoc builds.
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install a basic C toolchain (some rust crates compile C in their build.rs
+# scripts) and unzip (used to extract the pinned protoc release below).
 RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
        dnf install -y --setopt=install_weak_deps=False \
-            make unzip \
+            ca-certificates curl git gcc gcc-c++ make unzip \
        && dnf clean all && rm -rf /var/cache/dnf; \
    else \
        apt-get update -y \
        && apt-get install -y --no-install-recommends \
-            make unzip \
+            ca-certificates curl git build-essential unzip \
        && rm -rf /var/lib/apt/lists/*; \
    fi

 COPY tools/install_protoc.sh /tmp/install_protoc.sh
 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh

+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
 WORKDIR /workspace

-COPY requirements/build/rust.txt requirements/build/rust.txt
-RUN --mount=type=cache,target=/opt/uv/cache \
-    uv pip install --python /opt/venv/bin/python3 -r requirements/build/rust.txt
-
-# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
-# by the wheel build stage.
+# Copy only the rust workspace — the binary is the sole artifact we need.
 COPY rust rust
 COPY rust-toolchain.toml rust-toolchain.toml
-COPY tools/build_rust.py tools/build_rust.py
 COPY build_rust.sh build_rust.sh

 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4

-# Build the release artifacts. Cache cargo registry/git, but not target/,
-# because stale target metadata can outlive source updates across BuildKit
-# cache reuse.
+# Build the release binary. Cache cargo registry/git and target/, but copy the
+# binary out of the target/ cache mount so it persists into the image layer
+# for later COPY --from=rust-build.
 RUN --mount=type=cache,target=/root/.cargo/registry \
    --mount=type=cache,target=/root/.cargo/git \
-    bash build_rust.sh
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh
 #################### RUST BUILD IMAGE ####################

 #################### CSRC BUILD IMAGE ####################
@@ -340,7 +342,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 WORKDIR /workspace

 COPY pyproject.toml setup.py CMakeLists.txt ./
-COPY tools/build_rust.py tools/build_rust.py
 COPY cmake cmake/
 COPY csrc csrc/
 COPY vllm/envs.py vllm/envs.py
@@ -505,10 +506,9 @@ WORKDIR /workspace
 COPY --from=csrc-build /workspace/dist /precompiled-wheels
 COPY . .

-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
-COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs

 ARG GIT_REPO_CHECK=0
 RUN --mount=type=bind,source=.git,target=.git \
@@ -757,10 +757,10 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.12
+ARG FLASHINFER_VERSION=0.6.11.post2
 RUN --mount=type=cache,target=/opt/uv/cache \
    uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
-        --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
@@ -93,34 +93,35 @@ ENV DEBIAN_FRONTEND=noninteractive

 RUN apt-get update -y \
    && apt-get install -y --no-install-recommends \
-        ca-certificates curl git build-essential unzip python3 python3-pip \
+        ca-certificates curl git build-essential unzip \
    && rm -rf /var/lib/apt/lists/*

 COPY tools/install_protoc.sh /tmp/install_protoc.sh
 RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh

+# Install rustup; the toolchain itself is pinned by rust-toolchain.toml.
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
+    sh -s -- -y --profile minimal --default-toolchain none
+ENV PATH="/root/.cargo/bin:${PATH}"
+
 WORKDIR /workspace

-COPY requirements/build/rust.txt requirements/build/rust.txt
-RUN python3 -m pip install --no-cache-dir -r requirements/build/rust.txt
-
-# Copy only the Rust build inputs; build_rust.sh publishes artifacts needed
-# by the wheel build stage.
+# Copy only the rust workspace — the binary is the sole artifact we need.
 COPY rust rust
 COPY rust-toolchain.toml rust-toolchain.toml
-COPY tools/build_rust.py tools/build_rust.py
 COPY build_rust.sh build_rust.sh

 # Cap cargo parallelism to avoid exhausting the CI host's open-file limit
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4

-# Build the release artifacts. Cache cargo registry/git, but not target/,
-# because stale target metadata can outlive source updates across BuildKit
-# cache reuse.
-RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \
-    --mount=type=cache,target=/root/.cargo/git,sharing=locked \
-    bash build_rust.sh
+# Build the release binary. Cache cargo registry/git and target/, but copy the
+# binary out of the target/ cache mount so it persists into the image layer
+# for later COPY --from=rust-build.
+RUN --mount=type=cache,target=/root/.cargo/registry \
+    --mount=type=cache,target=/root/.cargo/git \
+    --mount=type=cache,target=/workspace/rust/target \
+    VLLM_RS_TARGET_PATH=/workspace/vllm-rs bash build_rust.sh

 ######################### BUILD IMAGE #########################
 FROM base AS vllm-build
@@ -153,10 +154,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \

 COPY . .

-# Drop the pre-built Rust artifacts into the source tree. setup.py detects
-# them and ships them as-is, skipping the local Rust build.
-COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs
-COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/
+# Drop the pre-built rust frontend binary into the source tree. setup.py
+# detects it and ships it as-is, skipping the local cargo build.
+COPY --from=rust-build /workspace/vllm-rs vllm/vllm-rs

 RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi

--- a/Show More
+++ b/Show More