Merge branch 'main' into wentao-optimize-model-runner-v2-sampler

Signed-off-by: yewentao256 <zhyanwentao@126.com>
[Refactor] Remove dead code (#42889 )
2026-05-18 22:45:59 +00:00 · 2026-05-18 16:41:22 -04:00 · 2026-05-18 19:41:58 +00:00 · 2026-05-18 11:56:22 -07:00 · 2026-05-18 11:22:56 -07:00 · 2026-05-18 11:14:21 -07:00
230 changed files with 12096 additions and 3501 deletions
@@ -12,7 +12,7 @@ steps:
  - vllm/_custom_ops.py
  - tests/kernels/attention/test_cpu_attn.py
  - tests/kernels/moe/test_cpu_fused_moe.py
-  - tests/kernels/moe/test_cpu_fp8_fused_moe.py
+  - tests/kernels/moe/test_cpu_quant_fused_moe.py
  - tests/kernels/test_onednn.py
  - tests/kernels/test_awq_int4_to_int8.py
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -21,7 +21,7 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-      pytest -x -v -s tests/kernels/moe/test_cpu_fp8_fused_moe.py
+      pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
@@ -70,7 +70,7 @@ steps:
  - tests/quantization/test_cpu_wna16.py
  commands:
    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
      
@@ -8,10 +8,3 @@ steps:
    commands: 
    - bash .buildkite/scripts/hardware_ci/run-hpu-test.sh

-  - label: "Intel GPU Test"
-    depends_on: []
-    soft_fail: true
-    device: intel_gpu
-    no_plugin: true
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -49,7 +49,7 @@ steps:
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
-      pytest -v -s lora/test_moe_lora_align_sum.py'
+      pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"'

 - label: LoRA Punica Kernels
  timeout_in_minutes: 45
@@ -7,6 +7,9 @@ env:
  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
  CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
  CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
+  MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
+  MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
+  MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"

 steps:
  - input: "Provide Release version here"
@@ -24,7 +27,7 @@ steps:
        agents:
          queue: arm64_cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
@@ -66,7 +69,7 @@ steps:
        agents:
          queue: cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
@@ -136,6 +139,8 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --progress plain \
@@ -162,6 +167,8 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --progress plain \
@@ -185,6 +192,8 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -210,6 +219,8 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -234,6 +245,8 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --progress plain \
@@ -261,6 +274,8 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --progress plain \
@@ -286,6 +301,8 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -312,6 +329,8 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
+              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
+              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -114,8 +114,7 @@ handle_pytest_exit() {
 # unquoted since they have no spaces and work fine.
 #
 # Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting values injected by
-# apply_rocm_test_overrides.
+# through untouched to avoid double-quoting well-formed shell fragments.
 #
 # NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
 # double-quotes stripped by the calling shell (see header comment).
@@ -247,102 +246,6 @@ re_quote_pytest_markers() {
  echo "${output% }"
 }

-###############################################################################
-# ROCm-specific pytest command rewrites
-#
-# These apply ignore flags and environment overrides for tests that are not
-# yet supported or behave differently on ROCm hardware. Kept as a single
-# function so new exclusions are easy to add in one place.
-###############################################################################
-
-apply_rocm_test_overrides() {
-  local cmds="$1"
-
-  # --- Model registry filter ---
-  if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
-    cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
-  fi
-
-  # --- LoRA: disable custom paged attention ---
-  if [[ $cmds == *"pytest -v -s lora"* ]]; then
-    cmds=${cmds//"pytest -v -s lora"/"pytest -v -s lora"}
-  fi
-
-  # --- Kernel ignores ---
-  if [[ $cmds == *" kernels/core"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/core/test_fused_quant_layernorm.py \
-    --ignore=kernels/core/test_permute_cols.py"
-  fi
-
-  if [[ $cmds == *" kernels/attention"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/attention/test_attention_selector.py \
-    --ignore=kernels/attention/test_encoder_decoder_attn.py \
-    --ignore=kernels/attention/test_flash_attn.py \
-    --ignore=kernels/attention/test_flashinfer.py \
-    --ignore=kernels/attention/test_prefix_prefill.py \
-    --ignore=kernels/attention/test_cascade_flash_attn.py \
-    --ignore=kernels/attention/test_mha_attn.py \
-    --ignore=kernels/attention/test_lightning_attn.py \
-    --ignore=kernels/attention/test_attention.py"
-  fi
-
-  if [[ $cmds == *" kernels/quantization"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/quantization/test_int8_quant.py \
-    --ignore=kernels/quantization/test_machete_mm.py \
-    --ignore=kernels/quantization/test_block_fp8.py \
-    --ignore=kernels/quantization/test_block_int8.py \
-    --ignore=kernels/quantization/test_marlin_gemm.py \
-    --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
-    --ignore=kernels/quantization/test_int8_kernel.py"
-  fi
-
-  if [[ $cmds == *" kernels/mamba"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/mamba/test_mamba_mixer2.py \
-    --ignore=kernels/mamba/test_causal_conv1d.py \
-    --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
-  fi
-
-  if [[ $cmds == *" kernels/moe"* ]]; then
-    cmds="${cmds} \
-    --ignore=kernels/moe/test_moe.py \
-    --ignore=kernels/moe/test_cutlass_moe.py"
-  fi
-
-  # --- Entrypoint ignores ---
-  if [[ $cmds == *" entrypoints/openai "* ]]; then
-    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
-    --ignore=entrypoints/openai/chat_completion/test_audio.py \
-    --ignore=entrypoints/openai/completion/test_shutdown.py \
-    --ignore=entrypoints/openai/test_completion.py \
-    --ignore=entrypoints/openai/models/test_models.py \
-    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
-    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
-    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
-  fi
-
-  if [[ $cmds == *" entrypoints/serve"* ]]; then
-    cmds="${cmds} \
-    --ignore=entrypoints/serve/lora/test_lora_adapters.py"
-  fi
-
-  if [[ $cmds == *" entrypoints/llm "* ]]; then
-    cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
-    --ignore=entrypoints/llm/test_chat.py \
-    --ignore=entrypoints/llm/test_accuracy.py \
-    --ignore=entrypoints/llm/test_init.py \
-    --ignore=entrypoints/llm/test_prompt_validation.py "}
-  fi
-
-  # Clean up escaped newlines from --ignore appends
-  cmds=$(echo "$cmds" | sed 's/ \\ / /g')
-
-  echo "$cmds"
-}
-
 ###############################################################################
 # Main
 ###############################################################################
@@ -409,7 +312,6 @@ else
  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
 fi

-commands=$(apply_rocm_test_overrides "$commands")
 echo "Final commands: $commands"

 MYPYTHONPATH=".."
@@ -1,53 +0,0 @@
-#!/bin/bash
-
-# This script build the CPU docker image and run the offline inference inside the container.
-# It serves a sanity check for compilation and basic model usage.
-set -ex
-
-image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
-container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
-
-# Try building the docker image
-docker build -t "${image_name}" -f docker/Dockerfile.xpu .
-
-# Setup cleanup
-remove_docker_container() {
-  docker rm -f "${container_name}" || true
-}
-trap remove_docker_container EXIT
-
-# Run the image and test offline inference/tensor parallel
-docker run \
-    --device /dev/dri:/dev/dri \
-    --net=host \
-    --ipc=host \
-    --privileged \
-    -v /dev/dri/by-path:/dev/dri/by-path \
-    --entrypoint="" \
-    -e "HF_TOKEN=${HF_TOKEN}" \
-    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK}" \
-    --name "${container_name}" \
-    "${image_name}" \
-    bash -c '
-    set -e
-    echo $ZE_AFFINITY_MASK
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager --max-model-len 8192
-    python3 examples/basic/offline_inference/generate.py --model 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 OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc --block-size 64 --enforce-eager --max-model-len 8192
-    cd tests
-    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
-    pytest -v -s v1/engine
-    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
-    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
-    pytest -v -s v1/structured_output
-    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
-    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py
-    pytest -v -s v1/test_serial_utils.py
-'
@@ -27,4 +27,5 @@ steps:
  - vllm/compilation
  commands:
    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
-    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
+    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
+    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
@@ -9,44 +9,7 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - |
-        if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
-          MAX_WAIT=300
-          INTERVAL=60
-          ELAPSED=0
-          while :; do
-            RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
-            HTTP_CODE=$(printf %s "$RAW" | tail -n1)
-            BODY=$(printf %s "$RAW" | head -n -1)
-            if [ "$HTTP_CODE" != "200" ]; then
-              echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
-              break
-            fi
-            STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
-            CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
-            if [ "$STATUS" = "none" ]; then
-              echo "no pre-run-check found for this commit; skipping gate."
-              break
-            fi
-            if [ -n "$CONCLUSION" ]; then
-              echo "pre-run-check conclusion: $CONCLUSION"
-              if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
-                echo "pre-run-check did not pass; failing docs build."
-                exit 1
-              fi
-              break
-            fi
-            if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
-              echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
-              break
-            fi
-            echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
-            sleep "$INTERVAL"
-            ELAPSED=$((ELAPSED + INTERVAL))
-          done
-        else
-          echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
-        fi
+      - bash docs/pre_run_check.sh
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -311,7 +311,6 @@ set(VLLM_EXT_SRC
  "csrc/attention/paged_attention_v1.cu"
  "csrc/attention/paged_attention_v2.cu"
  "csrc/attention/merge_attn_states.cu"
-  "csrc/attention/vertical_slash_index.cu"
  "csrc/pos_encoding_kernels.cu"
  "csrc/activation_kernels.cu"
  "csrc/layernorm_kernels.cu"
@@ -115,6 +115,39 @@ class SampleRequest:
 def sample_requests(
    tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
 ) -> list[SampleRequest]:
+    def _apply_random_prefix(
+        tokenizer: PreTrainedTokenizerBase,
+        requests: list[SampleRequest],
+        prefix_len: int,
+        seed: int,
+    ) -> list[SampleRequest]:
+        if prefix_len <= 0:
+            return requests
+        rng = np.random.default_rng(seed)
+        vocab_size = tokenizer.vocab_size
+        prohibited = getattr(tokenizer, "all_special_ids", None) or []
+        allowed = np.array([i for i in range(vocab_size) if i not in prohibited])
+        if len(allowed) == 0:
+            return requests
+        prefix_ids = rng.integers(0, len(allowed), size=prefix_len)
+        prefix_token_ids = allowed[prefix_ids].tolist()
+        out = []
+        for req in requests:
+            prompt_ids = tokenizer(req.prompt, add_special_tokens=False).input_ids
+            full_ids = prefix_token_ids + prompt_ids
+            full_prompt = tokenizer.decode(full_ids, skip_special_tokens=False)
+            out.append(
+                SampleRequest(
+                    prompt=full_prompt,
+                    prompt_len=len(tokenizer(full_prompt).input_ids),
+                    expected_output_len=req.expected_output_len,
+                    schema=req.schema,
+                    structure_type=req.structure_type,
+                    completion=req.completion,
+                )
+            )
+        return out
+
    if args.dataset == "json" or args.dataset == "json-unique":
        if args.json_schema_path is None:
            dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -261,6 +294,9 @@ def sample_requests(
                )
            )

+    requests = _apply_random_prefix(
+        tokenizer, requests, args.random_prefix_len, args.seed
+    )
    return requests


@@ -945,6 +981,15 @@ def create_argument_parser():
        "results in a more uniform arrival of requests.",
    )
    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--random-prefix-len",
+        type=int,
+        default=0,
+        help=(
+            "Number of prefix tokens to prepend to every prompt. "
+            "The same prefix is used for all prompts to enable prefix caching."
+        ),
+    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
@@ -1,401 +0,0 @@
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT license.
-
-#include <assert.h>
-
-#include <cuda.h>
-
-#include <torch/all.h>
-
-__device__ int64_t save_blocks(int* block_offset, int64_t range_start,
-                               int64_t range_end, int64_t block_size,
-                               int64_t input_block_count, int64_t kv_seqlen) {
-  if (range_start >= kv_seqlen) {
-    return input_block_count;
-  }
-  if (range_end > kv_seqlen) {
-    range_end = kv_seqlen;
-  }
-  int64_t current_block_count = input_block_count;
-  for (int idx = range_start; idx < range_end; idx += block_size) {
-    block_offset[current_block_count++] = idx;
-  }
-  return current_block_count;
-}
-
-__global__ void convert_vertical_slash_indexes_kernel(
-    const int* q_seqlens,         // [BATCH, ]
-    const int* kv_seqlens,        // [BATCH, ]
-    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
-    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
-    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
-    int64_t NNZ_V, int64_t NNZ_S,
-    bool causal  // True for intra, False for succ
-) {
-  const int batch_idx = blockIdx.y;
-  const int head_idx = blockIdx.x;
-  const int group_idx = blockIdx.z;
-
-  int64_t q_seqlen = q_seqlens[batch_idx];
-  int64_t kv_seqlen = kv_seqlens[batch_idx];
-  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
-  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
-  if (start_m >= q_seqlen) {
-    return;
-  }
-  int64_t end_m = start_m + BLOCK_SIZE_M;
-  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
-  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
-  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
-  block_count += row_offset;
-  block_offset += row_offset * NNZ_S;
-  column_count += row_offset;
-  column_index += row_offset * NNZ_V;
-
-  bool has_slash = true;
-  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
-  int64_t s = 0, v = 0;
-  int64_t v_idx = vertical_indexes[v++];
-  int64_t s_idx = slash_indexes[s++];
-  if (causal) {
-    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
-      s_idx = slash_indexes[s++];
-    }
-    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
-    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
-  } else {
-    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
-      s_idx = slash_indexes[s++];
-    }
-    if (s_idx > end_m + kv_seqlen) has_slash = false;
-    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
-  }
-
-  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
-  if (!has_slash) {
-    if (causal) {
-      range_start = (kv_seqlen - q_seqlen) + end_m;
-      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
-    } else {
-      range_start = kv_seqlen;
-      range_end = kv_seqlen + BLOCK_SIZE_N;
-    }
-  }
-
-  bool slash_finished = false;
-  while (1) {
-    if (v_idx < range_end) {
-      if (v_idx < range_start) {
-        column_index[tmp_col_cnt++] = v_idx;
-      }
-      if (v < NNZ_V) {
-        v_idx = vertical_indexes[v++];
-      } else {
-        if (causal)
-          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
-        else
-          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
-      }
-    } else {
-      if ((s < NNZ_S && causal) ||
-          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
-        if (causal)
-          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
-                      BLOCK_SIZE_M);
-        else
-          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
-      } else {
-        if (v == NNZ_V || (v_idx > range_start && causal)) {
-          // add the last vertical if no more slash
-          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
-            column_index[tmp_col_cnt++] = v_idx;
-          }
-          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-          break;
-        } else {
-          if (causal) {
-            range_start = (kv_seqlen - q_seqlen) + end_m;
-            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
-          } else {
-            // if slash_finished but there are vertical left, save current
-            // blocks
-            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-            range_start = kv_seqlen;
-            range_end = kv_seqlen + BLOCK_SIZE_N;
-          }
-          slash_finished = true;
-        }
-      }
-      if (!slash_finished) {
-        if (s_idx > range_end + BLOCK_SIZE_M) {
-          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-          range_start = s_idx - BLOCK_SIZE_M;
-          range_end = s_idx;
-        } else if (s_idx > range_end) {
-          range_end += BLOCK_SIZE_M;
-        }
-      }
-    }
-  }
-
-  block_count[0] = tmp_blk_cnt;
-  column_count[0] = tmp_col_cnt;
-}
-
-void convert_vertical_slash_indexes_64x64(
-    const int* q_seqlens,         // [BATCH, ]
-    const int* kv_seqlens,        // [BATCH, ]
-    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
-    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
-    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
-    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
-  const int N_THREADS = 64;
-  const dim3 dimBlock(N_THREADS);
-  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
-  convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock>>>(
-      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count,
-      block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M,
-      BLOCK_SIZE_N, NNZ_V, NNZ_S, causal);
-}
-
-/**
- * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
- *
- * This function builds the index of each row of blocks from vertical indices
- * and slash indices. The vertical indices are treated as points, while the
- * slash indices are converted as ranges. The output consists of the merged
- * ranges and separate column indices, where the ranges are represented by
- * block indices.
- *
- * The implementation is referenced from the original MInference repo:
- * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu.
- */
-void convert_vertical_slash_indexes(
-    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& block_offset,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
-    torch::Tensor& column_count,     // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& column_index,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
-    torch::Tensor q_seqlens,         // [BATCH, ]
-    torch::Tensor kv_seqlens,        // [BATCH, ]
-    torch::Tensor vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    torch::Tensor slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
-    bool causal) {
-  cudaSetDevice(q_seqlens.get_device());
-
-  int batch_size = slash_indexes.size(0);
-  int num_heads = slash_indexes.size(1);
-  int nnz_slash = slash_indexes.size(2);
-  int nnz_vertical = vertical_indexes.size(2);
-  int num_rows = (context_size + block_size_M - 1) / block_size_M;
-
-  convert_vertical_slash_indexes_64x64(
-      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
-      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
-      block_count.data_ptr<int>(), block_offset.data_ptr<int>(),
-      column_count.data_ptr<int>(), column_index.data_ptr<int>(), batch_size,
-      num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash,
-      causal);
-}
-
-__global__ void convert_vertical_slash_indexes_kernel_mergehead(
-    const int* q_seqlens,         // [BATCH, ]
-    const int* kv_seqlens,        // [BATCH, ]
-    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    const int* per_head_vertical_topkv, const int* per_head_slash_topkv,
-    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
-    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
-    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
-    int64_t NNZ_V, int64_t NNZ_S,
-    bool causal  // True for intra, False for succ
-) {
-  const int batch_idx = blockIdx.y;
-  const int head_idx = blockIdx.x;
-  const int group_idx = blockIdx.z;
-
-  int64_t q_seqlen = q_seqlens[batch_idx];
-  int64_t kv_seqlen = kv_seqlens[batch_idx];
-  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
-  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
-  if (start_m >= q_seqlen) {
-    return;
-  }
-  int64_t end_m = start_m + BLOCK_SIZE_M;
-  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
-  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
-  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
-  block_count += row_offset;
-  block_offset += row_offset * NNZ_S;
-  column_count += row_offset;
-  column_index += row_offset * NNZ_V;
-
-  // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
-  // above is buffer size, use to compute offset)
-  NNZ_S = per_head_slash_topkv[head_idx];
-  NNZ_V = per_head_vertical_topkv[head_idx];
-
-  bool has_slash = true;
-  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
-  int64_t s = 0, v = 0;
-  int64_t v_idx = vertical_indexes[v++];
-  int64_t s_idx = slash_indexes[s++];
-  if (causal) {
-    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
-      s_idx = slash_indexes[s++];
-    }
-    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
-    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
-  } else {
-    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
-      s_idx = slash_indexes[s++];
-    }
-    if (s_idx > end_m + kv_seqlen) has_slash = false;
-    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
-  }
-
-  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
-  if (!has_slash) {
-    if (causal) {
-      range_start = (kv_seqlen - q_seqlen) + end_m;
-      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
-    } else {
-      range_start = kv_seqlen;
-      range_end = kv_seqlen + BLOCK_SIZE_N;
-    }
-  }
-
-  bool slash_finished = false;
-  while (1) {
-    if (v_idx < range_end) {
-      if (v_idx < range_start) {
-        column_index[tmp_col_cnt++] = v_idx;
-      }
-      if (v < NNZ_V) {
-        v_idx = vertical_indexes[v++];
-      } else {
-        if (causal)
-          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
-        else
-          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
-      }
-    } else {
-      if ((s < NNZ_S && causal) ||
-          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
-        if (causal)
-          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
-                      BLOCK_SIZE_M);
-        else
-          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
-      } else {
-        if (v == NNZ_V || (v_idx > range_start && causal)) {
-          // add the last vertical if no more slash
-          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
-            column_index[tmp_col_cnt++] = v_idx;
-          }
-          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-          break;
-        } else {
-          if (causal) {
-            range_start = (kv_seqlen - q_seqlen) + end_m;
-            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
-          } else {
-            // if slash_finished but there are vertical left, save current
-            // blocks
-            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-            range_start = kv_seqlen;
-            range_end = kv_seqlen + BLOCK_SIZE_N;
-          }
-          slash_finished = true;
-        }
-      }
-      if (!slash_finished) {
-        if (s_idx > range_end + BLOCK_SIZE_M) {
-          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
-                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
-          range_start = s_idx - BLOCK_SIZE_M;
-          range_end = s_idx;
-        } else if (s_idx > range_end) {
-          range_end += BLOCK_SIZE_M;
-        }
-      }
-    }
-  }
-
-  block_count[0] = tmp_blk_cnt;
-  column_count[0] = tmp_col_cnt;
-}
-
-void convert_vertical_slash_indexes_64x64_mergehead(
-    const int* q_seqlens,         // [BATCH, ]
-    const int* kv_seqlens,        // [BATCH, ]
-    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    int* per_head_vertical_topkv, int* per_head_slash_topkv,
-    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
-    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
-    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
-    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
-    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
-  const int N_THREADS = 64;
-  const dim3 dimBlock(N_THREADS);
-  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
-  convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock>>>(
-      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes,
-      per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset,
-      column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N,
-      NNZ_V, NNZ_S, causal);
-}
-
-/**
- * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
- *
- * Like the above convert_vertical_slash_indexes, but with
- * pre-computed vertical and slash counts.
- */
-void convert_vertical_slash_indexes_mergehead(
-    torch::Tensor& block_count,            // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& block_offset,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
-    torch::Tensor& column_count,           // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& column_index,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
-    torch::Tensor q_seqlens,               // [BATCH, ]
-    torch::Tensor kv_seqlens,              // [BATCH, ]
-    torch::Tensor vertical_indexes,        // [BATCH, N_HEADS, NNZ_V]
-    torch::Tensor slash_indexes,           // [BATCH, N_HEADS, NNZ_S]
-    torch::Tensor vertical_indices_count,  // [N_HEADS, ]
-    torch::Tensor slash_indices_count,     // [N_HEADS, ]
-    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
-    bool causal) {
-  cudaSetDevice(q_seqlens.get_device());
-
-  int batch_size = slash_indexes.size(0);
-  int num_heads = slash_indexes.size(1);
-  int nnz_slash = slash_indexes.size(2);
-  int nnz_vertical = vertical_indexes.size(2);
-  int num_rows = (context_size + block_size_M - 1) / block_size_M;
-
-  convert_vertical_slash_indexes_64x64_mergehead(
-      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
-      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
-      vertical_indices_count.data_ptr<int>(),
-      slash_indices_count.data_ptr<int>(), block_count.data_ptr<int>(),
-      block_offset.data_ptr<int>(), column_count.data_ptr<int>(),
-      column_index.data_ptr<int>(), batch_size, num_heads, num_rows,
-      block_size_M, block_size_N, nnz_vertical, nnz_slash, causal);
-}
@@ -459,7 +459,7 @@ class AttentionScheduler {
    const int64_t kv_len_per_thread =
        (((total_kv_len / thread_num) + kv_len_alignment - 1) /
         kv_len_alignment) *
-        kv_len_alignment * (use_gqa ? input.num_heads_kv : input.num_heads_q);
+        kv_len_alignment;
    std::vector<AttentionWorkItemGroup> workitems;
    std::vector<ReductionWorkItemGroup> reduce_workitems;
    workitems.reserve(1024);
@@ -348,6 +348,13 @@ inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
  return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
 }

+template <>
+inline int get_cache_blocks<uint8_t>(int chunk_size) {
+  // mxfp4 uses bf16 as accumulate type
+  int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
+  return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
+}
+
 // 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
 template <typename T, typename func_t>
 inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
@@ -213,7 +213,7 @@ struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
      weight + nb_start * width,                                                             \
      out + bs * seqlen * dim + mb_start * dim + nb_start,                                   \
      has_bias ? bias + nb_start : nullptr,                                                  \
-      has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \
+      has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
      has_initial_states_value,                                                              \
      mb_size,                                                                               \
      dim,                                                                                   \
@@ -233,7 +233,8 @@ void causal_conv1d_fwd_kernel_impl(
    int64_t dim,
    int64_t seqlen,
    int64_t width,
-    int64_t num_seq_blocks) {
+    int64_t num_seq_blocks,
+    int64_t conv_state_slot_stride) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -282,7 +283,7 @@ void causal_conv1d_fwd_kernel_impl(
    at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
      for (int64_t bs = begin; bs < end; ++bs) {
        update_conv_state(
-            conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
+            conv_states + bs * conv_state_slot_stride, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
      }
    });
  }
@@ -316,7 +317,8 @@ void causal_conv1d_fwd_varlen_kernel_impl(
    int64_t batch,
    int64_t dim,
    int64_t width,
-    int64_t num_seq_blocks) {
+    int64_t num_seq_blocks,
+    int64_t conv_state_slot_stride) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -366,7 +368,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
        int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
        int32_t batch_offset = query_start_loc[bs];
        update_conv_state(
-            conv_states + conv_state_index * (width - 1) * dim,
+            conv_states + conv_state_index * conv_state_slot_stride,
            input + batch_offset * dim,
            width,
            dim,
@@ -389,7 +391,8 @@ void causal_conv1d_update_kernel_impl(
    int64_t batch,
    int64_t dim,
    int64_t seqlen,
-    int64_t width) {
+    int64_t width,
+    int64_t conv_state_slot_stride) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -430,7 +433,7 @@ void causal_conv1d_update_kernel_impl(
    });
  });

-#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim
+#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*conv_state_slot_stride + (w) * dim

  // update conv_states
  at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
@@ -592,6 +595,9 @@ at::Tensor causal_conv1d_fwd_cpu(
    }
  }

+// IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
+  int64_t conv_state_slot_stride = conv_states->stride(0);
+
  // block size for sequence blocks, 32
  constexpr int64_t BLOCK_M = block_size_m();

@@ -618,7 +624,8 @@ at::Tensor causal_conv1d_fwd_cpu(
          batch,
          dim,
          width,
-          num_seq_blocks);
+          num_seq_blocks,
+          conv_state_slot_stride);
    } else {
      causal_conv1d_fwd_kernel_impl<scalar_t>(
          out.data_ptr<scalar_t>(),
@@ -633,7 +640,8 @@ at::Tensor causal_conv1d_fwd_cpu(
          dim,
          seqlen,
          width,
-          num_seq_blocks);
+          num_seq_blocks,
+          conv_state_slot_stride);
    }
  });
  return out;
@@ -690,6 +698,8 @@ at::Tensor causal_conv1d_update_cpu(
    conv_states.copy_(conv_states_copy);
  }

+  // IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
+  int64_t conv_state_slot_stride = conv_states.stride(0);
  at::Tensor out = at::empty_like(x);
  AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
    causal_conv1d_update_kernel_impl<scalar_t>(
@@ -703,7 +713,8 @@ at::Tensor causal_conv1d_update_cpu(
        batch,
        dim,
        seqlen,
-        width);
+        width,
+        conv_state_slot_stride);
  });
  return out;
 }
@@ -847,6 +847,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
    int64_t v_strideB,
    int64_t v_strideS,
    int64_t v_strideH,
+    int64_t state_slot_stride,
    bool use_qk_l2norm_in_kernel,
    double softplus_threshold) {
  using bVec = at::vec::Vectorized<scalar_t>;
@@ -907,7 +908,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
    data_index_init(begin, bi, batch_size, si, seq_len, ni, v_num_heads);
    for (int64_t i = begin; i < end; ++i) {
      int64_t cache_index = indices_ptr[bi];
-      int64_t state_offset = (cache_index * v_num_heads + ni) * head_dim * v_head_dim;
+      int64_t state_offset = cache_index * state_slot_stride + ni * head_dim * v_head_dim;
      float g_val = -std::exp(float(A_log_ptr[ni])) *
                    softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold);
      float g_val_exp = std::exp(g_val);
@@ -1321,6 +1322,8 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
  int64_t v_strideB = v.stride(1);
  int64_t v_strideS = v.stride(0);
  int64_t v_strideH = v.stride(2);
+  // IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
+  int64_t state_slot_stride = initial_state_source.stride(0);
  at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
  at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);

@@ -1353,6 +1356,7 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
            v_strideB,
            v_strideS,
            v_strideH,
+            state_slot_stride,
            use_qk_l2norm_in_kernel,
            softplus_threshold);
      });
@@ -72,7 +72,17 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
  return use_int8_w8a8 ? K + sizeof(int32_t) : K;
 }

-enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };
+enum class CPUAcTMethod : int { silu_and_mul = 0, swiglu = 1 };
+
+constexpr bool operator==(CPUAcTMethod a, int b) {
+  return static_cast<int>(a) == b;
+}
+
+constexpr bool operator==(int a, CPUAcTMethod b) {
+  return a == static_cast<int>(b);
+}
+
+enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };

 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
  return static_cast<int64_t>(a) == b;
@@ -99,6 +109,9 @@ inline int64_t get_4bit_block_k_size(int64_t group_size) {
 // pack weight to vnni format
 at::Tensor convert_weight_packed(at::Tensor& weight);

+// pack scale to blocked format for mxfp4
+at::Tensor convert_scale_packed(at::Tensor& scale);
+
 // pack weight to vnni format for int4
 std::tuple<at::Tensor, at::Tensor, at::Tensor>
 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
@@ -129,9 +142,9 @@ void fused_experts_int8_kernel_impl(
    int64_t topk,
    int64_t num_tokens_post_pad);

-// moe implementations for fp8 w8a16
-template <typename scalar_t>
-void fused_experts_fp8_kernel_impl(
+// moe implementations for fp8 w8a16 and mxfp4
+template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
+void fused_experts_fp_kernel_impl(
    scalar_t* __restrict__ output,
    scalar_t* __restrict__ ic0,
    scalar_t* __restrict__ ic1,
@@ -140,10 +153,12 @@ void fused_experts_fp8_kernel_impl(
    scalar_t* __restrict__ B_tmp,
    float* __restrict__ C_tmp,
    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
+    const packed_t* __restrict__ packed_w1,
+    const packed_t* __restrict__ packed_w2,
+    const float* __restrict__ w1_bias,
+    const float* __restrict__ w2_bias,
+    const param_t* __restrict__ w1s,
+    const param_t* __restrict__ w2s,
    int64_t block_size_N,
    int64_t block_size_K,
    const float* __restrict__ topk_weights,
@@ -155,7 +170,11 @@ void fused_experts_fp8_kernel_impl(
    int64_t K,
    int64_t E,
    int64_t topk,
-    int64_t num_tokens_post_pad);
+    int64_t num_tokens_post_pad,
+    float alpha,
+    float limit,
+    CPUAcTMethod act_func,
+    bool with_bias);

 // shared expert implementation for int8 w8a8
 template <typename scalar_t>
@@ -266,6 +285,7 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
+    const float* __restrict__ Bbias,
    const float* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -294,6 +314,26 @@ void tinygemm_kernel(
    int64_t ldc,
    bool brg);

+// mxfp4
+template <typename scalar_t>
+void tinygemm_kernel(
+    const scalar_t* __restrict__ A,
+    const uint8_t* __restrict__ B,
+    scalar_t* __restrict__ C,
+    scalar_t* __restrict__ Btmp,
+    float* __restrict__ Ctmp,
+    const float* __restrict__ Bbias,
+    const uint8_t* __restrict__ scale,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg,
+    int64_t block_size_K,
+    bool do_unpack = true);
+
 template <typename scalar_t>
 void tinygemm_kernel(
    scalar_t* C,
@@ -67,6 +67,23 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
  }
 }

+template <>
+inline void
+copy_add_stub(float* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = fVec::size();
+
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    fVec data = fVec::loadu(input + d) + fVec::loadu(bias + d);
+    data.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = input[d] + bias[d];
+  }
+}
+
 inline void unpack_B(
    at::BFloat16* __restrict__ Btmp,
    const at::Float8_e4m3fn* __restrict__ packed_B,
@@ -343,7 +360,6 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, float, has_bias, BLOC
    Unroll<ROWS * COLS>{}(storec);
  }
 };
-
 template <int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
  static inline void apply(
@@ -918,6 +934,7 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
+    const float* __restrict__ Bbias,
    const float* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -928,6 +945,11 @@ void tinygemm_kernel(
    bool brg,
    int64_t block_size_K,
    bool do_unpack) {
+  if (Bbias != nullptr) {
+    tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
+        A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+    return;
+  }
  tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
      A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
 }
@@ -949,7 +971,6 @@ void tinygemm_kernel(
    bool brg) {
  tinygemm_kernel2<scalar_t>(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg);
 }
-
 template <typename scalar_t>
 void tinygemm_kernel(
    const scalar_t* __restrict__ A,
@@ -957,6 +978,7 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
+    const float* __restrict__ Bbias,
    const uint8_t* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -967,10 +989,68 @@ void tinygemm_kernel(
    bool brg,
    int64_t block_size_K,
    bool do_unpack) {
+  if (Bbias != nullptr) {
+    tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
+        A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+    return;
+  }
  tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
      A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
 }

+// tinygemm interface
+template <typename scalar_t>
+void tinygemm_kernel(
+    const scalar_t* __restrict__ A,
+    const at::Float8_e4m3fn* __restrict__ B,
+    float* __restrict__ C,
+    scalar_t* __restrict__ Btmp,
+    const float* __restrict__ Bbias,
+    const float* __restrict__ scale,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg,
+    int64_t block_size_K,
+    bool do_unpack) {
+  if (Bbias != nullptr) {
+    tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
+        A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+    return;
+  }
+  tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
+      A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+}
+
+template <typename scalar_t>
+void tinygemm_kernel(
+    const scalar_t* __restrict__ A,
+    const uint8_t* __restrict__ B,
+    float* __restrict__ C,
+    scalar_t* __restrict__ Btmp,
+    const float* __restrict__ Bbias,
+    const uint8_t* __restrict__ scale,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg,
+    int64_t block_size_K,
+    bool do_unpack) {
+  if (Bbias != nullptr) {
+    tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
+        A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+    return;
+  }
+  tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
+      A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
+}
+
 #define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
  template void tinygemm_kernel<TYPE_A>(                      \
      const TYPE_A* __restrict__ A,                           \
@@ -978,6 +1058,7 @@ void tinygemm_kernel(
      TYPE_A* __restrict__ C,                                 \
      TYPE_A* __restrict__ Btmp,                              \
      float* __restrict__ Ctmp,                               \
+      const float* __restrict__ Bbias,                        \
      const TYPE_S* __restrict__ scale,                       \
      int64_t M,                                              \
      int64_t N,                                              \
@@ -1020,7 +1101,6 @@ inline const float* get_bias_data(const std::optional<at::Tensor>& bias, int64_t
  }
  return nullptr;
 }
-
 // FP8 and MXFP4 WoQ uses the same pattern:
 //   Btmp : [T, BLOCK_N * K]
 //   Ctmp : [T, BLOCK_M * BLOCK_N]
@@ -812,6 +812,7 @@ void shared_expert_kernel_impl(
 static inline void check_moe_scales(
    bool use_int8_w8a8,
    bool use_fp8_w8a16,
+    bool use_mxfp4,
    const std::optional<at::Tensor>& w1_scale,
    const std::optional<at::Tensor>& w2_scale,
    const std::optional<std::vector<int64_t>> block_size) {
@@ -825,6 +826,12 @@ static inline void check_moe_scales(
    TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
    TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
  }
+  if (use_mxfp4) {
+    TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4.");
+    TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
+    TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
+    TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8.");
+  }
 }

 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1)                      \
@@ -839,8 +846,8 @@ static inline void check_moe_scales(
  TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))

 // hidden_states: [M, K]
-// w1: [E, 2N, K]
-// w2: [E, K, N]
+// w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
+// w2: [E, K, N] or [E, K, N / 2] for uint8
 // topk_weights: [M, topk]
 // topk_ids: [M, topk] (int32_t)
 //
@@ -858,6 +865,10 @@ at::Tensor fused_experts_cpu(
    const std::optional<at::Tensor>& w1_zero,
    const std::optional<at::Tensor>& w2_zero,
    const std::optional<std::vector<int64_t>> block_size,
+    const std::optional<at::Tensor>& w1_bias,
+    const std::optional<at::Tensor>& w2_bias,
+    const std::optional<double>& alpha,
+    const std::optional<double>& limit,
    bool is_vnni) {
  auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
  auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -897,8 +908,12 @@ at::Tensor fused_experts_cpu(
  int64_t topk = topk_weights_.size(1);

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

  // check weight shapes
  CHECK_EQ(w2.size(0), E);
@@ -911,6 +926,7 @@ at::Tensor fused_experts_cpu(
  check_moe_scales(
      moe_comp_method == CPUQuantMethod::INT8_W8A8,
      moe_comp_method == CPUQuantMethod::FP8_W8A16,
+      moe_comp_method == CPUQuantMethod::MXFP4,
      w1_scale,
      w2_scale,
      block_size);
@@ -965,7 +981,7 @@ at::Tensor fused_experts_cpu(
  //   5. Aq_tmp : [M, K] or [M * topk, N]
  //   6. As_tmp : [M * topk]
  //
-  // for fp8 w8a16:
+  // for fp8 w8a16 and mxfp4:
  //   7. intermediate_cache0 : [M * topk, 2N]
  //   8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
  //
@@ -978,7 +994,7 @@ at::Tensor fused_experts_cpu(
  if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
    buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
  }
-  if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
+  if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) {
    buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
  }
  if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
@@ -1031,9 +1047,11 @@ at::Tensor fused_experts_cpu(
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
      scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
      scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
+      bool with_bias = w1_bias.has_value();
+      auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;

      CHECK_MOE_SCALES_FP8(1, 2);
-      fused_experts_fp8_kernel_impl(
+      fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
          out_hidden_states.data_ptr<scalar_t>(),
          intermediate_cache0,
          intermediate_cache1,
@@ -1044,6 +1062,8 @@ at::Tensor fused_experts_cpu(
          hidden_states.data_ptr<scalar_t>(),
          packed_w1.data_ptr<at::Float8_e4m3fn>(),
          packed_w2.data_ptr<at::Float8_e4m3fn>(),
+          with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
+          with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
          w1s.data_ptr<float>(),
          w2s.data_ptr<float>(),
          block_size_N,
@@ -1057,7 +1077,56 @@ at::Tensor fused_experts_cpu(
          K,
          E,
          topk,
-          num_tokens_post_pad);
+          num_tokens_post_pad,
+          alpha.has_value() ? float(alpha.value()) : 0,
+          limit.has_value() ? float(limit.value()) : 0,
+          act_func,
+          with_bias);
+    } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
+      scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
+      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
+      scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
+      scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
+      bool with_bias = w1_bias.has_value();
+      auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
+
+      // mxfp4 supports only group size of 32 (2^5)
+      constexpr int64_t group_size = 32;
+      auto w1s = w1_scale.value();
+      auto w2s = w2_scale.value();
+      TORCH_CHECK(w1s.numel() == E * 2 * N * K / group_size, "w1_scale size mismatch");
+      TORCH_CHECK(w2s.numel() == E * K * N / group_size, "w2_scale size mismatch");
+      fused_experts_fp_kernel_impl<scalar_t, uint8_t, uint8_t, true>(
+          out_hidden_states.data_ptr<scalar_t>(),
+          intermediate_cache0,
+          intermediate_cache1,
+          intermediate_cache2,
+          A_tmp,
+          B_tmp,
+          C_tmp,
+          hidden_states.data_ptr<scalar_t>(),
+          packed_w1.data_ptr<uint8_t>(),
+          packed_w2.data_ptr<uint8_t>(),
+          with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
+          with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
+          w1s.data_ptr<uint8_t>(),
+          w2s.data_ptr<uint8_t>(),
+          /*block_size_N*/ 1,
+          /*block_size_K*/ group_size,
+          topk_weights_.data_ptr<float>(),
+          sorted_ids,
+          expert_ids,
+          offsets,
+          M,
+          N,
+          K,
+          E,
+          topk,
+          num_tokens_post_pad,
+          alpha.has_value() ? float(alpha.value()) : 0,
+          limit.has_value() ? float(limit.value()) : 0,
+          act_func,
+          with_bias);
    } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
      uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1185,7 +1254,7 @@ at::Tensor shared_expert_cpu(
  CHECK_EQ(packed_w2.size(1), packed_N);

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

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

@@ -176,3 +176,109 @@ inline void silu_and_mul_stub(
    out_vec.store(out + d);
  }
 }
+
+template <typename scalar_t>
+inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  const fVec weight_vec = fVec(weight);
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    fVec data0 = fVec::loadu(input + d) * weight_vec;
+    fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
+    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
+    out_vec.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<scalar_t>(input[d] * weight);
+  }
+}
+
+// input = input + input2
+inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = fVec::size();
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    fVec x_fvec = fVec::loadu(input + d);
+    fVec y_fvec = fVec::loadu(input2 + d);
+    x_fvec = x_fvec + y_fvec;
+    x_fvec.store(input + d);
+  }
+  for (; d < size; ++d) {
+    input[d] = input[d] + input2[d];
+  }
+}
+
+template <typename scalar_t>
+inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  const fVec weight_vec = fVec(weight);
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    bVec x = bVec::loadu(input + d);
+    fVec x0, x1;
+    std::tie(x0, x1) = at::vec::convert_to_float(x);
+    x0 = x0 * weight_vec;
+    x1 = x1 * weight_vec;
+    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<scalar_t>(input[d] * weight);
+  }
+}
+
+template <typename scalar_t>
+inline void clamp_sigmoid_and_mul_stub(
+    scalar_t* __restrict__ out,
+    const scalar_t* __restrict__ input,
+    int64_t size,
+    const float alpha,
+    const float limit) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  const fVec one = fVec(1.f);
+  const fVec zero = fVec(0.f);
+  const fVec limit_v = fVec(limit);
+  const fVec nlimit_v = fVec(-limit);
+  const fVec alpha_v = fVec(alpha);
+
+  // no remainder
+#pragma GCC unroll 4
+  for (int64_t d = 0; d < size; d += bVec::size()) {
+    bVec x = bVec::loadu(input + d);
+    fVec x0_, y0_;
+    std::tie(x0_, y0_) = at::vec::convert_to_float(x);
+    float tmp_buffer[fVec::size() * 2];  // 32
+    float tmp_glu[fVec::size()];         // 16
+    float tmp_linear[fVec::size()];      // 16
+    x0_.store(tmp_buffer);
+    y0_.store(tmp_buffer + fVec::size());
+    // interleaved: x[2i] = glu, x[2i+1] = linear
+    for (int j = 0; j < fVec::size(); ++j) {
+      // x0 [0,2,..30]
+      tmp_glu[j] = tmp_buffer[j * 2];
+      // y0 [1,3,...31]
+      tmp_linear[j] = tmp_buffer[j * 2 + 1];
+    }
+    fVec x0 = fVec::loadu(tmp_glu);
+    fVec y0 = fVec::loadu(tmp_linear);
+
+    // clamp
+    x0 = at::vec::minimum(x0, limit_v);
+    y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
+    // x * sigmoid(x * alpha)
+    x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
+    // (y + 1) * x
+    y0 = y0 + one;
+    x0 = x0 * y0;
+    convert_from_float_and_store<scalar_t>(out + d / 2, x0);
+  }
+}
@@ -7,8 +7,8 @@
 #include "gemm.h"
 #include "moe.h"

-template <typename scalar_t>
-void fused_experts_fp8_kernel_impl(
+template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
+void fused_experts_fp_kernel_impl(
    scalar_t* __restrict__ output,
    scalar_t* __restrict__ ic0,
    scalar_t* __restrict__ ic1,
@@ -17,10 +17,12 @@ void fused_experts_fp8_kernel_impl(
    scalar_t* __restrict__ B_tmp,
    float* __restrict__ C_tmp,
    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
+    const packed_t* __restrict__ packed_w1,
+    const packed_t* __restrict__ packed_w2,
+    const float* __restrict__ w1_bias,
+    const float* __restrict__ w2_bias,
+    const param_t* __restrict__ w1s,
+    const param_t* __restrict__ w2s,
    int64_t block_size_N,
    int64_t block_size_K,
    const float* __restrict__ topk_weights,
@@ -32,7 +34,11 @@ void fused_experts_fp8_kernel_impl(
    int64_t K,
    int64_t E,
    int64_t topk,
-    int64_t num_tokens_post_pad) {
+    int64_t num_tokens_post_pad,
+    float alpha,
+    float limit,
+    CPUAcTMethod act_func,
+    bool with_bias) {
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();

@@ -42,12 +48,20 @@ void fused_experts_fp8_kernel_impl(
  int64_t scale_size_N = div_up(2 * N, block_size_N);
  int64_t scale_size_K = div_up(K, block_size_K);
  int64_t blocks_n_per_group = block_size_N / BLOCK_N;
+  std::function<int64_t(int64_t)> scale_offset_per_block;
+  if constexpr (is_mxfp4) {
+    scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
+  } else {
+    scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
+  }

-  const int64_t stride_e = 2 * N * K;
-  const int64_t stride_n = K;
+  const int64_t packed_K = get_row_size<packed_t>(K);
+
+  const int64_t stride_e = 2 * N * packed_K;
+  const int64_t stride_n = packed_K;

  int64_t avg_M = std::max(int64_t(1), M * topk / E);
-  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
+  const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);

  int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);

@@ -57,14 +71,15 @@ void fused_experts_fp8_kernel_impl(
    int tid = get_thread_num();
    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;

-    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);

      // B shape [K, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
-      const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
-      const float* __restrict__ Bs =
-          w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+      const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
+      const param_t* __restrict__ Bs =
+          w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
+      const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;

      // do unpacking for the first row or a new expert
      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -88,6 +103,7 @@ void fused_experts_fp8_kernel_impl(
          /*   C            */ ic0 + offset * 2 * N + nb * BLOCK_N,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
+          /*   Bbias        */ B_bias,
          /*   scale        */ Bs,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -106,12 +122,20 @@ void fused_experts_fp8_kernel_impl(
  });

  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
-  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
-    }
-  });
-
+  if (act_func == CPUAcTMethod::silu_and_mul) {
+    at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+      for (int64_t m = begin; m < end; ++m) {
+        silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
+      }
+    });
+  } else if (act_func == CPUAcTMethod::swiglu) {
+    at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+      for (int64_t m = begin; m < end; ++m) {
+        clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
+        clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
+      }
+    });
+  }
  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
  //   w2 : [E, K, N] as [E, OC, IC]
  const int64_t OC = K;  // rename K as OC
@@ -120,15 +144,16 @@ void fused_experts_fp8_kernel_impl(
  const int64_t NB2 = div_up(OC, BLOCK_N);
  scale_size_N = div_up(K, block_size_N);
  scale_size_K = div_up(N, block_size_K);
-  const int64_t stride_e2 = OC * IC;
-  const int64_t stride_oc = IC;
+  const int64_t packed_IC = get_row_size<packed_t>(IC);
+  const int64_t stride_e2 = OC * packed_IC;
+  const int64_t stride_oc = packed_IC;

  // parallel on [MB2, NB2]
  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    int tid = get_thread_num();
    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];

-    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = offsets[mb + 1] - offsets[mb];
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

@@ -139,9 +164,10 @@ void fused_experts_fp8_kernel_impl(

      // B shape [IC, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
-      const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
-      const float* __restrict__ Bs =
-          w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+      const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
+      const param_t* __restrict__ Bs =
+          w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
+      const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;

      // do unpacking for the first row or a new expert
      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -153,6 +179,7 @@ void fused_experts_fp8_kernel_impl(
          /*   C            */ C,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
+          /*   Bbias        */ B_bias,
          /*   scale        */ Bs,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -177,7 +204,6 @@ void fused_experts_fp8_kernel_impl(
      at::native::cpublas::brgemm_release();
    }
  });
-
  // stage 3: out = intermediate_cache2.sum(dim=1)
  //   from [M, topk, K] to [M, K]
  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
@@ -187,35 +213,43 @@ void fused_experts_fp8_kernel_impl(
  });
 }

-#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE)             \
-  template void fused_experts_fp8_kernel_impl<TYPE>(   \
-      TYPE* __restrict__ output,                       \
-      TYPE* __restrict__ ic0,                          \
-      TYPE* __restrict__ ic1,                          \
-      TYPE* __restrict__ ic2,                          \
-      TYPE* __restrict__ A_tmp,                        \
-      TYPE* __restrict__ B_tmp,                        \
-      float* __restrict__ C_tmp,                       \
-      const TYPE* __restrict__ input,                  \
-      const at::Float8_e4m3fn* __restrict__ packed_w1, \
-      const at::Float8_e4m3fn* __restrict__ packed_w2, \
-      const float* __restrict__ w1s,                   \
-      const float* __restrict__ w2s,                   \
-      int64_t block_size_N,                            \
-      int64_t block_size_K,                            \
-      const float* __restrict__ topk_weights,          \
-      const int32_t* __restrict__ sorted_ids,          \
-      const int32_t* __restrict__ expert_ids,          \
-      const int32_t* __restrict__ offsets,             \
-      int64_t M,                                       \
-      int64_t N,                                       \
-      int64_t K,                                       \
-      int64_t E,                                       \
-      int64_t topk,                                    \
-      int64_t num_tokens_post_pad)
+#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4)           \
+  template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
+      TYPE1* __restrict__ output,                                            \
+      TYPE1* __restrict__ ic0,                                               \
+      TYPE1* __restrict__ ic1,                                               \
+      TYPE1* __restrict__ ic2,                                               \
+      TYPE1* __restrict__ A_tmp,                                             \
+      TYPE1* __restrict__ B_tmp,                                             \
+      float* __restrict__ C_tmp,                                             \
+      const TYPE1* __restrict__ input,                                       \
+      const TYPE2* __restrict__ packed_w1,                                   \
+      const TYPE2* __restrict__ packed_w2,                                   \
+      const float* __restrict__ w1_bias,                                     \
+      const float* __restrict__ w2_bias,                                     \
+      const TYPE3* __restrict__ w1s,                                         \
+      const TYPE3* __restrict__ w2s,                                         \
+      int64_t block_size_N,                                                  \
+      int64_t block_size_K,                                                  \
+      const float* __restrict__ topk_weights,                                \
+      const int32_t* __restrict__ sorted_ids,                                \
+      const int32_t* __restrict__ expert_ids,                                \
+      const int32_t* __restrict__ offsets,                                   \
+      int64_t M,                                                             \
+      int64_t N,                                                             \
+      int64_t K,                                                             \
+      int64_t E,                                                             \
+      int64_t topk,                                                          \
+      int64_t num_tokens_post_pad,                                           \
+      float alpha,                                                           \
+      float limit,                                                           \
+      CPUAcTMethod act_func,                                                 \
+      bool with_bias)

-INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
-INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);
+INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
+INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
+INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
+INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);

 template <typename scalar_t>
 void shared_expert_fp8_kernel_impl(
@@ -266,6 +300,7 @@ void shared_expert_fp8_kernel_impl(
          /*   C            */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
+          /*   Bbias        */ nullptr,
          /*   scale        */ w1s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -317,6 +352,7 @@ void shared_expert_fp8_kernel_impl(
          /*   C            */ C,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
+          /*   Bbias        */ nullptr,
          /*   scale        */ w2s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -11,7 +11,7 @@

 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
-
+#include <immintrin.h>
 namespace {

 using namespace at::vec;
@@ -21,6 +21,15 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
  return at::vec::convert_from_float<scalar_t>(a, b);
 }

+template <typename scalar_t>
+inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
+  float out_buffer[at::vec::Vectorized<float>::size()];
+  a.store(out_buffer);
+  for (int i = 0; i < 16; i++) {
+    out[i] = (scalar_t)out_buffer[i];
+  }
+}
+
 // allow f16, bf16
 template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
 inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
@@ -50,6 +59,11 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
  return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
 }

+template <>
+inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
+  _mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
+}
+
 #define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))

 #define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
@@ -125,7 +139,6 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
  return cvt_e4m3_bf16_intrinsic_with_denorm(a);
 #endif
 }
-
 // faster version of float8_e4m3fn conversion to bfloat16
 //
 // we mapped cuda implementation from below link and vectorized with avx512:
@@ -64,6 +64,8 @@ at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,

 at::Tensor convert_weight_packed(at::Tensor& weight);

+at::Tensor convert_scale_packed(at::Tensor& scale);
+
 at::Tensor fused_experts_cpu(
    at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
    at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
@@ -71,7 +73,11 @@ at::Tensor fused_experts_cpu(
    const std::optional<at::Tensor>& w2_scale,
    const std::optional<at::Tensor>& w1_zero,
    const std::optional<at::Tensor>& w2_zero,
-    const std::optional<std::vector<int64_t>> block_size, bool is_vnni);
+    const std::optional<std::vector<int64_t>> block_size,
+    const std::optional<at::Tensor>& w1_bias,
+    const std::optional<at::Tensor>& w2_bias,
+    const std::optional<double>& alpha, const std::optional<double>& limit,
+    bool is_vnni);

 at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
                                     at::Tensor& scales2,
@@ -98,6 +104,45 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                              at::Tensor& w_scales,
                              std::optional<at::Tensor> bias);

+// Adapted from sglang: GDN
+std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
+    const at::Tensor& query, const at::Tensor& key, const at::Tensor& value,
+    const at::Tensor& g, const at::Tensor& beta,
+    const at::Tensor& initial_state, bool output_final_state,
+    const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel,
+    double eps = 1e-5);
+
+at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
+    const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
+    const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,
+    const at::Tensor& b, at::Tensor& initial_state_source,
+    const at::Tensor& initial_state_indices, const at::Tensor& cu_seqlens,
+    bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,
+    double softplus_threshold = 20.0);
+
+std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(
+    const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,
+    const at::Tensor& dt_bias);
+
+// Adapted from sglang: casual_conv1d kernels
+at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight);
+
+at::Tensor causal_conv1d_fwd_cpu(
+    const at::Tensor& x, const at::Tensor& weight,
+    const std::optional<at::Tensor>& bias,
+    const std::optional<at::Tensor>& conv_states,
+    const std::optional<at::Tensor>& query_start_loc,
+    const std::optional<at::Tensor>& cache_indices,
+    const std::optional<at::Tensor>& has_initial_state, bool silu_activation,
+    int64_t pad_slot_id, bool is_vnni);
+
+at::Tensor causal_conv1d_update_cpu(
+    const at::Tensor& x, const at::Tensor& conv_states,
+    const at::Tensor& weight, const std::optional<at::Tensor>& bias,
+    bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
+    const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
+    bool is_vnni);
+
 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                         const std::string& activation);

@@ -365,11 +410,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
  ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
  ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
+  ops.def("convert_scale_packed(Tensor! scale) -> Tensor");
+  ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
  ops.def(
      "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
      "topk_weights, Tensor topk_ids, bool "
      "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
-      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
+      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
+      "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
+      "bool is_vnni) -> "
      "Tensor");
  ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
  ops.def(
@@ -397,6 +446,47 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
      "bool is_vnni) -> Tensor");
  ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
+
+  // Adapted from sglang: GDN kernels
+  ops.def(
+      "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
+      "Tensor g, Tensor beta, "
+      "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool "
+      "head_first, "
+      "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");
+  ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU,
+           &chunk_gated_delta_rule_cpu);
+  ops.def(
+      "fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor "
+      "dt_bias, Tensor q, Tensor k, Tensor v, Tensor "
+      "a, Tensor b, Tensor(a!) initial_state_source, Tensor "
+      "initial_state_indices, Tensor cu_seqlens, bool "
+      "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "
+      "softplus_threshold=20.0) -> Tensor");
+  ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,
+           &fused_sigmoid_gating_delta_rule_update_cpu);
+  ops.def(
+      "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "
+      "-> (Tensor, Tensor)");
+  ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
+
+  // Adapted from sglang: casual_conv1d kernels
+  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
+  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
+           &causal_conv1d_weight_pack);
+  ops.def(
+      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
+      "conv_states, Tensor? query_start_loc,"
+      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
+      "int pad_slot_id, bool is_vnni) -> "
+      "Tensor");
+  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
+  ops.def(
+      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
+      "weight, Tensor? bias, bool silu_activation,"
+      "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
+      "bool is_vnni) -> Tensor");
+  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
 #endif

  // CPU attention kernels
@@ -202,7 +202,7 @@ struct cutlass_3x_gemm_sm120 {
              sizeof(typename CollectiveEpilogue::SharedStorage))>,
          KernelSchedule>::CollectiveOp;

-  using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -72,7 +72,7 @@ struct cutlass_3x_gemm_sm120_custom {
              sizeof(typename CollectiveEpilogue::SharedStorage))>,
          KernelSchedule, void>::CollectiveOp;

-  using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -15,11 +15,11 @@ SUPPORT_SM80 = False
 for arch in sys.argv[1].split(","):
    arch = arch[: arch.index(".") + 2].replace(".", "")
    arch = int(arch)
-    # only SM89 and SM120 fully support
-    # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
+    # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
+    # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
    # SM90 and SM100 can use this PTX, but it’s simulated
    # with FP16 MMA, so it cannot achieve any acceleration.
-    if arch in [89, 120]:
+    if arch == 89 or arch // 10 == 12:
        SUPPORT_FP8 = True
    if arch >= 80:
        SUPPORT_SM80 = True
@@ -448,8 +448,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                "FP8 only support Ada Lovelace or newer GPUs.");
    TORCH_CHECK(
        major_capability * 10 + minor_capability == 89 ||
-            major_capability * 10 + minor_capability == 120,
-        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
+            major_capability == 12,
+        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
        "Marlin W4A16 on other devices).");
  }

@@ -390,7 +390,13 @@ __global__ void moe_lora_align_block_size_kernel(
    int32_t* __restrict__ token_mask, bool has_expert_map) {
  int lora_idx = blockIdx.x / 2;
  int lora_id = lora_ids[lora_idx];
-  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
+  // Output buffers are indexed by lora_id (in [0, max_loras)). The grid
+  // iterates one extra slot to accommodate the "-1" entry that
+  // active_lora_ids may hold in position 0 for mixed base + LoRA batches;
+  // guard against any other unexpected lora_id >= max_loras to avoid
+  // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in
+  // the Triton _fused_moe_lora_kernel.
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
    return;
  }

@@ -420,10 +426,21 @@ __global__ void lora_count_and_sort_expert_tokens_kernel(
    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
    int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
    int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask,
-    int32_t* lora_ids, bool has_expert_map) {
+    int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled,
+    bool has_expert_map) {
  int lora_idx = blockIdx.x;
  int lora_id = lora_ids[lora_idx];
-  if (lora_id == -1) {
+  // Same guard rationale as moe_lora_align_block_size_kernel. Additionally
+  // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns
+  // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is
+  // allocated with torch::empty), so running the sort loop here would traverse
+  // garbage mask bits and pollute this slot's rows of sorted_token_ids and
+  // cumsum_buffer. Downstream consumers already skip disabled slots, so the
+  // pollution is dormant today, but the check keeps behavior symmetric with
+  // the other two align kernels and avoids O(numel) wasted work per disabled
+  // slot. Short-circuit evaluation ensures adapter_enabled is only indexed
+  // after lora_id is confirmed to be in [0, max_loras).
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
    return;
  }

@@ -446,7 +463,8 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
    int32_t* token_mask, bool has_expert_map) {
  int lora_idx = blockIdx.x;
  int lora_id = lora_ids[lora_idx];
-  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
+  // Same guard rationale as moe_lora_align_block_size_kernel.
+  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
    return;
  }

@@ -698,7 +716,15 @@ void moe_lora_align_block_size(
                  scalar_t, fill_threads>;
          AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
              (void*)kernel, shared_mem));
-          kernel<<<max_loras, blockDim, shared_mem, stream>>>(
+          // Grid size is (max_loras + 1) because active_lora_ids has length
+          // max_loras + 1: sorted-unique values of token_lora_mapping, which
+          // can include -1 (base-model tokens) in addition to up to max_loras
+          // real LoRA slots. Using max_loras would drop the real LoRA slot
+          // when -1 is present at position 0 and leave output buffers
+          // uninitialized, causing illegal memory accesses in downstream
+          // MoE-LoRA kernels. This mirrors the fix made for the Triton
+          // _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
+          kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
              topk_ids.data_ptr<scalar_t>(),
              token_lora_mapping.data_ptr<int32_t>(), block_size,
              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -722,10 +748,17 @@ void moe_lora_align_block_size(
          auto align_kernel =
              vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;

-          // launch two threadblocks for each lora
+          // Launch two threadblocks per LoRA slot, across max_loras + 1 slots
+          // to cover the extra "-1" (base-model tokens) entry that
+          // active_lora_ids may contain in addition to up to max_loras real
+          // LoRA slots. Using max_loras would drop the real LoRA slot when -1
+          // occupies position 0 and leave the output buffers uninitialized,
+          // causing illegal memory accesses downstream. Mirrors the grid fix
+          // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277.
          // blockIdx.x % 2 == 0: counting experts and aligning
          // blockIdx.x % 2 == 1: filling sorted_token_ids
-          align_kernel<<<max_loras * 2, blockDim, shared_mem_size, stream>>>(
+          align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
+                         stream>>>(
              topk_ids.data_ptr<scalar_t>(),
              token_lora_mapping.data_ptr<int32_t>(), block_size,
              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -744,7 +777,10 @@ void moe_lora_align_block_size(
          const int max_blocks = 65535;
          const int actual_blocks = std::min(num_blocks, max_blocks);

-          dim3 gridDims(max_loras, actual_blocks);
+          // Same rationale as align_kernel above: iterate over max_loras + 1
+          // slots so the sort kernel processes the real LoRA slot even when
+          // active_lora_ids has -1 at position 0.
+          dim3 gridDims(max_loras + 1, actual_blocks);
          auto sort_kernel =
              vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;

@@ -753,7 +789,8 @@ void moe_lora_align_block_size(
              sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
              expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
              max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
-              lora_ids.data_ptr<int32_t>(), has_expert_map);
+              max_loras, lora_ids.data_ptr<int32_t>(),
+              adapter_enabled.data_ptr<int32_t>(), has_expert_map);
        }
      });
 }
@@ -92,16 +92,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bool use_fp32_reduce, bool is_zp_float,"
      "int thread_k, int thread_n, int blocks_per_sm) -> Tensor");

-  m.def(
-      "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
-      "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
-      "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, "
-      "int b_q_type, SymInt size_m, "
-      "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int "
-      "topk, "
-      "int moe_block_size, bool replicate_input, bool apply_weights)"
-      " -> Tensor");
-
  m.def(
      "moe_permute(Tensor input, Tensor topk_ids,"
      "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
@@ -60,32 +60,6 @@ void merge_attn_states(
    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
    const std::optional<int64_t> prefill_tokens_with_context,
    const std::optional<torch::Tensor>& output_scale = std::nullopt);
-#ifndef USE_ROCM
-void convert_vertical_slash_indexes(
-    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& block_offset,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
-    torch::Tensor& column_count,     // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& column_index,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
-    torch::Tensor q_seqlens,         // [BATCH, ]
-    torch::Tensor kv_seqlens,        // [BATCH, ]
-    torch::Tensor vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
-    torch::Tensor slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
-    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
-    bool causal);
-
-void convert_vertical_slash_indexes_mergehead(
-    torch::Tensor& block_count,            // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& block_offset,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
-    torch::Tensor& column_count,           // [BATCH, N_HEADS, NUM_ROWS]
-    torch::Tensor& column_index,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
-    torch::Tensor q_seqlens,               // [BATCH, ]
-    torch::Tensor kv_seqlens,              // [BATCH, ]
-    torch::Tensor vertical_indexes,        // [BATCH, N_HEADS, NNZ_V]
-    torch::Tensor slash_indexes,           // [BATCH, N_HEADS, NNZ_S]
-    torch::Tensor vertical_indices_count,  // [N_HEADS, ]
-    torch::Tensor slash_indices_count, int64_t context_size,
-    int64_t block_size_M, int64_t block_size_N, bool causal);
-#endif

 void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
              double epsilon);
@@ -15,11 +15,11 @@ SUPPORT_SM80 = False
 for arch in sys.argv[1].split(","):
    arch = arch[: arch.index(".") + 2].replace(".", "")
    arch = int(arch)
-    # only SM89 and SM120 fully support
-    # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
+    # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
+    # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
    # SM90 and SM100 can use this PTX, but it’s simulated
    # with FP16 MMA, so it cannot achieve any acceleration.
-    if arch in [89, 120]:
+    if arch == 89 or arch // 10 == 12:
        SUPPORT_FP8 = True
    if arch >= 80:
        SUPPORT_SM80 = True
@@ -76,30 +76,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "    int!? prefill_tokens_with_context,"
      "    Tensor? output_scale=None) -> ()");
  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
-#ifndef USE_ROCM
-  ops.def(
-      "convert_vertical_slash_indexes("
-      "   Tensor! block_count, Tensor! block_offset, "
-      "   Tensor! column_count, Tensor! column_index, "
-      "   Tensor q_seqlens, Tensor q_seqlens, "
-      "   Tensor vertical_indexes, Tensor slash_indexes, "
-      "   int context_size, int block_size_M, int block_size_N, "
-      "   bool causal) -> ()");
-  ops.impl("convert_vertical_slash_indexes", torch::kCUDA,
-           &convert_vertical_slash_indexes);
-
-  ops.def(
-      "convert_vertical_slash_indexes_mergehead("
-      "   Tensor! block_count, Tensor! block_offset, "
-      "   Tensor! column_count, Tensor! column_index, "
-      "   Tensor q_seqlens, Tensor q_seqlens, "
-      "   Tensor vertical_indexes, Tensor slash_indexes, "
-      "   Tensor vertical_indices_count, Tensor slash_indices_count, "
-      "   int context_size, int block_size_M, int block_size_N, "
-      "   bool causal) -> ()");
-  ops.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA,
-           &convert_vertical_slash_indexes_mergehead);
-#endif

  // Activation ops
  // Activation function used in SwiGLU.
@@ -639,7 +639,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # 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.8.post1
+ARG FLASHINFER_VERSION=0.6.11.post2
 RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -858,6 +858,28 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
    fi

+# Optional override: install mooncake-transfer-engine from a URL instead of the
+# PyPI release pulled in above. Use this for wheels built with non-default CMake
+# flags (e.g. `STORE_USE_ETCD=ON` for master HA). The URL's manylinux glibc
+# floor must be <= the FINAL_BASE_IMAGE's glibc.
+ARG MOONCAKE_WHEEL_AARCH64
+ARG MOONCAKE_WHEEL_X86_64
+RUN if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
+        if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+            WHEEL="${MOONCAKE_WHEEL_AARCH64}"; \
+        else \
+            WHEEL="${MOONCAKE_WHEEL_X86_64}"; \
+        fi && \
+        if [ -n "${WHEEL}" ]; then \
+            uv pip install --system "${WHEEL}" && \
+            CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
+            if [ ! -f /usr/local/cuda/lib64/libcudart.so ] && \
+               [ -f "/usr/local/cuda/lib64/libcudart.so.${CUDA_MAJOR}" ]; then \
+                ln -s "libcudart.so.${CUDA_MAJOR}" /usr/local/cuda/lib64/libcudart.so; \
+            fi; \
+        fi; \
+    fi
+
 ENV VLLM_USAGE_SOURCE production-docker-image
 ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
    VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
@@ -217,13 +217,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.


 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.8.post1
+# release version: v0.6.11.post2
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
    --mount=type=cache,target=/root/.cache/uv \
    echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.8.post1 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
    && cd flashinfer \
    && git submodule update --init --recursive \
    && echo "finish git clone flashinfer..." \
@@ -68,7 +68,7 @@
      "default": "true"
    },
    "FLASHINFER_VERSION": {
-      "default": "0.6.8.post1"
+      "default": "0.6.11.post2"
    },
    "GDRCOPY_CUDA_VERSION": {
      "default": "12.8"
@@ -34,6 +34,8 @@ th {
 | HuggingFace-AIMO | ✅ | ✅ | `AI-MO/aimo-validation-aime`, `AI-MO/NuminaMath-1.5`, `AI-MO/NuminaMath-CoT` |
 | HuggingFace-Other | ✅ | ✅ | `lmms-lab/LLaVA-OneVision-Data`, `Aeala/ShareGPT_Vicuna_unfiltered` |
 | HuggingFace-MTBench | ✅ | ✅ | `philschmid/mt-bench` |
+| HuggingFace-HumanEval | ✅ | ✅ | `openai/openai_humaneval` |
+| HuggingFace-GSM8K | ✅ | ✅ | `openai/gsm8k` |
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
 | HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
@@ -443,6 +445,26 @@ vllm bench serve \
    --num-prompts 80
 ```

+`openai/openai_humaneval`:
+
+``` bash
+vllm bench serve \
+    --model NousResearch/Hermes-3-Llama-3.1-8B \
+    --dataset-name hf \
+    --dataset-path openai/openai_humaneval \
+    --num-prompts 80
+```
+
+`openai/gsm8k`:
+
+``` bash
+vllm bench serve \
+    --model NousResearch/Hermes-3-Llama-3.1-8B \
+    --dataset-name hf \
+    --dataset-path openai/gsm8k \
+    --num-prompts 80
+```
+
 `vdaita/edit_5k_char` or `vdaita/edit_10k_char`:

 ``` bash
@@ -77,6 +77,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
 * `get_input_modality(...)` - return the modality of the inputs.
 * `get_max_frames_per_video()` - return model-specific max frames per video.
+* `postprocess_encoder_output(...)` - post process encoder output, directly call scatter_output_slices by default

 !!! note
    The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -89,6 +90,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
 | `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
+| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |

 !!! note
    Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
@@ -248,6 +248,57 @@ Now, you can specify a base_model_name alongside the name and path using JSON fo

 To provide the backward compatibility support, you can still use the old key-value format (name=path), but the `base_model_name` will remain unspecified in that case.

+## Mixing 2D and 3D MoE LoRA Adapters
+
+To serve 2D-format(based on `megatron`) and 3D-format (based on `peft`) adapters from the same engine instance, start the server with `--enable-mixed-moe-lora-format`
+and declare the layout of each adapter explicitly via the `is_3d_lora_weight` field.
+
+Server startup (static modules):
+
+```bash
+vllm serve Qwen/Qwen3.6-35B-A3B \
+    --enable-lora \
+    --enable-mixed-moe-lora-format \
+    --tensor-parallel-size 4 \
+    --enable-expert-parallel \
+    --lora-modules \
+        '{"name": "lora-2d", "path": "jeeejeee/qwen36-35ba3b-2d-weights-poken-lora", "is_3d_lora_weight": false}' \
+        '{"name": "lora-3d", "path": "jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora", "is_3d_lora_weight": true}'
+```
+
+Dynamic load via `/v1/load_lora_adapter`:
+
+```bash
+curl -X POST http://localhost:8000/v1/load_lora_adapter \
+-H "Content-Type: application/json" \
+-d '{
+    "lora_name": "lora-3d",
+    "lora_path": "/path/to/3d-format-lora",
+    "is_3d_lora_weight": true
+}'
+```
+
+!!! warning "You must know your adapter's layout"
+    Under `--enable-mixed-moe-lora-format`, vLLM trusts whatever
+    `is_3d_lora_weight` the caller declares — it does **not** inspect the
+    checkpoint to verify. A wrong declaration will load weights into the
+    wrong stacked buffers and silently produce garbage outputs, with no
+    error at load time. Confirm the layout before serving:
+
+    - **2D (per-expert, megatron-style)** → set `is_3d_lora_weight: false`.
+      Adapter keys look like `...experts.{idx}.gate_proj.lora_A.weight`,
+      `...experts.{idx}.up_proj.lora_A.weight`,
+      `...experts.{idx}.down_proj.lora_A.weight` — one set per expert.
+    - **3D (fused, peft-style)** → set `is_3d_lora_weight: true`.
+      Adapter keys look like `...experts.gate_up_proj.lora_A.weight`,
+      `...experts.down_proj.lora_A.weight` — a single tensor that stacks
+      all experts on the leading dim.
+
+When `--enable-mixed-moe-lora-format` is **not** set, `is_3d_lora_weight`
+is ignored: vLLM picks the wrapper from the base model's
+`is_3d_moe_weight` and the adapter is required to match. The field is
+also ignored for non-MoE models.
+
 ## LoRA model lineage in model card

 The new format of `--lora-modules` is mainly to support the display of parent model information in the model card. Here's an explanation of how your current response supports this:
@@ -2,7 +2,7 @@

 MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:

- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
+- **CPU/disk offloading**: Extend effective KV cache capacity by offloading to CPU memory or disk via Mooncake's transfer engine.
 - **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
 - **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.

@@ -38,18 +38,29 @@ Create a JSON configuration file (e.g., `mooncake_config.json`):

 ```json
 {
+  "mode": "embedded",
  "metadata_server": "P2PHANDSHAKE",
  "master_server_address": "127.0.0.1:50051",
  "global_segment_size": "80GB",
  "local_buffer_size": "4GB",
  "protocol": "rdma",
-  "device_name": ""
+  "device_name": "",
+  "enable_offload": false
 }
 ```

+- `mode`: Topology selection. `"embedded"` (default, PR-40900 baseline) has each
+  vLLM rank contribute `global_segment_size` to the pool in-process.
+  `"standalone-store"` makes ranks pure requesters — an external
+  `mooncake_client` process owns the CPU pool and (optionally) the SSD tier.
 - `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
+- `global_segment_size`: CPU memory contributed to the distributed pool (per
+  GPU). Must be `> 0` in `embedded` mode and `0` in `standalone-store` mode.
 - `local_buffer_size`: Private buffer for this node's own operations (per GPU).
+- `enable_offload`: When `true`, vLLM allocates a DirectIO staging buffer so
+  large prefills do not exceed the owner's SSD-write budget. Set this together
+  with the matching `--enable_offload=true` flag on `mooncake_master` and on
+  the external `mooncake_client` (if any).

 Set the config path via environment variable:

@@ -126,12 +137,60 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \

 A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.

+### Disk Offloading
+
+Disk offloading is most commonly run in `standalone-store` mode: an external
+`mooncake_client` process owns the CPU pool and the SSD tier, and each vLLM
+rank is a pure requester. This avoids per-rank duplication of the SSD pool
+and keeps DirectIO budget tracking on a single process.
+
+Three things need to be aligned for end-to-end disk offloading:
+
+1. **`mooncake_master`** is started with `--enable_offload=true`.
+2. **`mooncake_client`** (the owner) is started with `--enable_offload=true`
+   plus an SSD path via `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`.
+3. **vLLM-side** sets `"enable_offload": true` in the JSON config file (this is
+   read by the connector and is **not** an environment variable).
+
+Example `mooncake_config.json` for the vLLM side:
+
+```json
+{
+  "mode": "standalone-store",
+  "metadata_server": "P2PHANDSHAKE",
+  "master_server_address": "127.0.0.1:50051",
+  "global_segment_size": 0,
+  "local_buffer_size": "4GB",
+  "protocol": "rdma",
+  "device_name": "mlx5_0",
+  "enable_offload": true
+}
+```
+
+Steer this rank to the local owner segment with:
+
+```bash
+export MOONCAKE_PREFERRED_SEGMENT=127.0.0.1:50053
+```
+
+The owner's SSD directory, on-disk eviction policy, and the DirectIO staging
+buffer size are controlled on the `mooncake_client` side via the standard
+Mooncake environment variables (`MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`,
+`MOONCAKE_BUCKET_EVICTION_POLICY`, `MOONCAKE_USE_URING`,
+`MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES`,
+`MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES`, etc.). Those are independent of
+the vLLM JSON config.
+
 ## Environment Variables

 | Variable | Description | Default |
 | --- | --- | --- |
 | `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
 | `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
+| `MOONCAKE_PREFERRED_SEGMENT` | Pin this rank's replicas to a specific owner segment (`host:port`); used in `standalone-store` mode | — |
+| `MOONCAKE_REQUESTER_LOCAL_HOSTNAME` | Override the hostname the vLLM rank registers with Mooncake as a requester. Defaults to the rank's resolved IP. | — |
+| `VLLM_MOONCAKE_STORE_TIER_LOG` | When `1`, logs a per-batch tier summary (memory vs disk hits) for observability | disabled |
+| `VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO` | Fraction of the owner's DirectIO staging buffer that the requester will fill in a single `batch_get_into_multi_buffers` call. Lower → more conservative pre-split, more round trips. | 0.9 |

 ## KV Transfer Config

@@ -464,6 +464,17 @@ Supported models:

 Flags: `--tool-call-parser gigachat3`

+### Apertus Models (`apertus`)
+
+Use the chat template from the examples folder; it fixes several OpenAI compatibility issues: `--chat-template /vllm-workspace/examples/tool_chat_template_apertus.jinja`
+
+Supported models:
+
+* `swiss-ai/Apertus-8B-Instruct-2509`
+* `swiss-ai/Apertus-70B-Instruct-2509`
+
+Flags: `--tool-call-parser apertus`
+
 ### Models with Pythonic Tool Calls (`pythonic`)

 A growing number of models output a python list to represent tool calls instead of using JSON. This has the advantage of inherently supporting parallel tool calls and removing ambiguity around the JSON schema required for tool calls. The `pythonic` tool parser can support such models.
@@ -0,0 +1,41 @@
+if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
+  MAX_WAIT=300
+  INTERVAL=60
+  ELAPSED=0
+  while :; do
+    RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
+    HTTP_CODE=$(printf %s "$RAW" | tail -n1)
+    BODY=$(printf %s "$RAW" | sed '$d')
+    if [ "$HTTP_CODE" != "200" ]; then
+      echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
+      break
+    fi
+    STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
+    CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
+    CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")")
+    if [ "$STATUS" = "none" ]; then
+      echo "no pre-run-check found for this commit; skipping gate."
+      break
+    fi
+    if [ -n "$CONCLUSION" ]; then
+      echo "pre-run-check conclusion: $CONCLUSION"
+      if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
+        echo "pre-run-check did not pass; skipping docs build."
+        if [ -n "$CHECK_URL" ]; then
+          echo "pre-run-check failure reason: $CHECK_URL"
+        fi
+        exit 1
+      fi
+      break
+    fi
+    if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
+      echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
+      break
+    fi
+    echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
+    sleep "$INTERVAL"
+    ELAPSED=$((ELAPSED + INTERVAL))
+  done
+else
+  echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
+fi
@@ -1,21 +1,37 @@
 # IPC Engine

-The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same node and same GPU**. This avoids any data copying, making it a efficient option when colocating training and inference.
+The IPC weight transfer engine uses **CUDA IPC** (Inter-Process Communication) handles to share GPU memory directly between the trainer and inference workers on the **same GPU**. This avoids any data copying, making it the most efficient option when colocating training and inference. Multi-GPU setups are supported — weights are all gathered by each GPU and are extracted by the correct colocated process.

 ## When to Use IPC

- Training and inference on the **same GPU** (colocated)
- You want to minimize memory overhead by sharing tensors in-place
+- Training and inference share the **same GPU(s)** (colocated)

 ## How It Works

-1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`.
-2. IPC handles are sent to the inference engine via **Ray.remote()** or **HTTP POST**.
-3. The inference worker reconstructs the tensors from the handles, reading directly from the trainer's GPU memory.
+1. The trainer creates CUDA tensors for each weight and generates IPC handles using `torch.multiprocessing.reductions.reduce_tensor`. In multi-GPU setups (e.g. FSDP), each trainer rank must all-gather the full tensor for each layer onto its own GPU before generating the IPC handle.
+2. IPC handles for each gpu are sent to the inference engine via **Ray**, **HTTP**, or a **custom callable**. Each rank only reads the handle corresponding to its own GPU.
+3. The inference worker reconstructs the tensors from the handles using `rebuild_cuda_tensor`, reading directly from the trainer's GPU memory.

 !!! warning
    IPC handles involve sending serialized Python objects. When using HTTP transport, you must set `VLLM_ALLOW_INSECURE_SERIALIZATION=1` on both the server and client. This is because IPC handles are pickled and base64-encoded for HTTP transmission.

+## Packed (Chunked) Transfer
+
+By default, all weights are sent in a single API call. For large models, this requires the full model to reside in GPU memory on both sides simultaneously. Setting `packed=True` enables **chunked transfer** with bounded GPU memory:
+
+- Weights are concatenated into fixed-size packed buffers (controlled by `packed_buffer_size_bytes`).
+- Each chunk is sent as a separate `update_weights` call within a single `start_weight_update` / `finish_weight_update` bracket, so the layerwise reload pass is initialized once at the start and finalized once at the end regardless of chunk count.
+- After each chunk is consumed, the GPU memory for that chunk can be reclaimed.
+
+```python
+trainer_args = IPCTrainerSendWeightsArgs(
+    send_mode="ray",
+    llm_handle=llm_actor_handle,
+    packed=True,
+    packed_buffer_size_bytes=256 * 1024 * 1024,  # 256 MB chunks
+)
+```
+
 ## Initialization

 The IPC backend requires no initialization on either side. The `init_transfer_engine` call is a no-op for IPC.
@@ -35,7 +51,7 @@ from vllm.distributed.weight_transfer.ipc_engine import (
 )

 trainer_args = IPCTrainerSendWeightsArgs(
-    mode="ray",
+    send_mode="ray",
    llm_handle=llm_actor_handle,
 )
 # start
@@ -57,7 +73,7 @@ Used when vLLM is running as an HTTP server:

 ```python
 trainer_args = IPCTrainerSendWeightsArgs(
-    mode="http",
+    send_mode="http",
    url="http://localhost:8000",
 )

@@ -77,7 +93,22 @@ response = requests.post(url, json={}, timeout=60)
 response.raise_for_status()
 ```

-In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. As with Ray mode, you must call `start_weight_update` before and `finish_weight_update` after.
+In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. Because the worker deserializes the payload via `pickle.loads`, the vLLM server must be started with `VLLM_ALLOW_INSECURE_SERIALIZATION=1`.
+
+```python
+def my_custom_sender(update_info: IPCWeightTransferUpdateInfo):
+    # Custom logic to deliver update_info to vLLM
+    ...
+
+trainer_args = IPCTrainerSendWeightsArgs(
+    send_mode=my_custom_sender,
+)
+
+IPCWeightTransferEngine.trainer_send_weights(
+    iterator=model.named_parameters(),
+    trainer_args=trainer_args,
+)
+```

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

@@ -128,7 +128,7 @@ firewall configuration instructions.

 ### Overview

-The `--api-key` flag (or `VLLM_API_KEY` environment variable) provides authentication for vLLM's HTTP server, but **only for OpenAI-compatible API endpoints under the `/v1` path prefix**. Many other sensitive endpoints are exposed on the same HTTP server without any authentication enforcement.
+The `--api-key` flag (or `VLLM_API_KEY` environment variable) provides authentication for vLLM's HTTP server, but **only for OpenAI-compatible API endpoints under the `/v1` path prefix**, and other similar `/v2`, `/inference` path prefix**. Many other sensitive endpoints are exposed on the same HTTP server without any authentication enforcement.

 **Important:** Do not rely exclusively on `--api-key` for securing access to vLLM. Additional security measures are required for production deployments.

@@ -154,6 +154,9 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok
 - `/v1/rerank` - Reranking API
 - `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
 - `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
+- `/inference/v1/generate` - Generate completions
+- `/v2/embed` - Cohere Embed API
+- `/v2/rerank` - Cohere Rerank API

 ### Unprotected Endpoints (No API Key Required)

@@ -162,7 +165,6 @@ The following endpoints **do not require authentication** even when `--api-key`
 **Inference endpoints:**

 - `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints)
- `/inference/v1/generate` - Generate completions
 - `/generative_scoring` - Generative scoring API
 - `/pooling` - Pooling API
 - `/classify` - Classification API
@@ -2560,6 +2560,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
    "qwen2_vl",
    "qwen3_5",
    "qwen3_5_moe",
+    "stepvl",
 ]


@@ -173,7 +173,7 @@ def main():
    start_weight_update(BASE_URL, is_checkpoint_format=False)

    print("Broadcasting weights via CUDA IPC (HTTP)...")
-    trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=BASE_URL)
+    trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
    IPCWeightTransferEngine.trainer_send_weights(
        iterator=train_model.named_parameters(),
        trainer_args=trainer_args,
@@ -70,10 +70,14 @@ class TrainModel:
            self.llm_handle.init_weight_transfer_engine.remote(dict(init_info=dict()))
        )

-    def broadcast_weights(self, llm_handle: ray.actor.ActorHandle):
+    def broadcast_weights(
+        self, llm_handle: ray.actor.ActorHandle, packed: bool = False
+    ):
        """Broadcast weights to the inference engine using IPC."""
        self.llm_handle = llm_handle
-        trainer_args = IPCTrainerSendWeightsArgs(mode="ray", llm_handle=llm_handle)
+        trainer_args = IPCTrainerSendWeightsArgs(
+            send_mode="ray", llm_handle=llm_handle, packed=packed
+        )
        IPCWeightTransferEngine.trainer_send_weights(
            iterator=self.train_model.named_parameters(),
            trainer_args=trainer_args,
@@ -141,10 +145,10 @@ ray.get(llm.finish_weight_update.remote())

 ray.get(llm.wake_up.remote(tags=["scheduling"]))

-# Generate text with the updated model.
-outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+outputs_packed = ray.get(llm.generate.remote(prompts, sampling_params))
 print("-" * 50)
-for output in outputs_updated:
+print("Results after packed/chunked IPC weight sync:")
+for output in outputs_packed:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
@@ -0,0 +1,425 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
+weight transfer and **packed** tensors.
+
+Layout (4 GPUs, TP=1, DP=4, EP):
+  * One Ray placement group per GPU.
+  * Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
+    (sync API) using fractional GPUs so both fit on the same device.
+  * The 4 ``LLM`` instances form a DP group via env-var-based SPMD
+    coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
+    mechanism used by ``examples/offline_inference/data_parallel.py``.
+  * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
+    waits for initialization, and orchestrates generation / weight-sync.
+
+Uses the built-in ``ray`` send_mode: each FSDP worker calls
+``trainer_send_weights`` targeting its colocated LLM actor.
+
+This example was run on 4xH100.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import asdict
+
+import ray
+import torch
+import torch.distributed as dist
+from huggingface_hub import snapshot_download
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from torch.distributed._tensor import DTensor
+from torch.distributed.fsdp import fully_shard
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.ipc_engine import (
+    IPCTrainerSendWeightsArgs,
+    IPCWeightTransferEngine,
+    IPCWeightTransferInitInfo,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+
+TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
+VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))
+
+MODEL_NAME = "Qwen/Qwen3-30B-A3B"
+
+FSDP_WORLD_SIZE = 4
+INFERENCE_TP_SIZE = 1
+INFERENCE_DP_SIZE = 4
+
+
+class MyLLM(LLM):
+    """LLM subclass that configures DP env vars for SPMD coordination."""
+
+    def __init__(
+        self,
+        *args,
+        dp_rank: int = 0,
+        dp_size: int = 1,
+        dp_master_ip: str = "127.0.0.1",
+        dp_master_port: int = 0,
+        **kwargs,
+    ):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
+        os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
+        os.environ["VLLM_DP_RANK"] = str(dp_rank)
+        os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
+        os.environ["VLLM_DP_SIZE"] = str(dp_size)
+        os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
+        os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
+
+        super().__init__(*args, **kwargs)
+
+    def ready(self):
+        return True
+
+
+@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
+class FSDPTrainWorker:
+    """One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""
+
+    def __init__(
+        self,
+        model_name: str,
+        rank: int,
+        fsdp_world_size: int,
+        fsdp_master_addr: str,
+        fsdp_master_port: int,
+    ):
+        self.rank = rank
+
+        os.environ["MASTER_ADDR"] = fsdp_master_addr
+        os.environ["MASTER_PORT"] = str(fsdp_master_port)
+
+        dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
+        torch.accelerator.set_device_index(0)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16
+        )
+
+        self.weight_names = [n for n, _ in model.named_parameters()]
+        self.weight_dtype_names = [
+            str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
+        ]
+        self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
+
+        for layer in model.model.layers:
+            fully_shard(layer)
+        fully_shard(model)
+
+        self.model = model
+
+    def get_rank(self):
+        return self.rank
+
+    def get_weight_metadata(self):
+        return self.weight_names, self.weight_dtype_names, self.weight_shapes
+
+    def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
+        """All-gather full params; all ranks create IPC handles, rank 0 sends.
+
+        All ranks must call trainer_send_weights so they participate in the
+        all_gather_object collective inside _all_gather_and_merge_handles.
+        Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
+        """
+
+        def _full_param_iter():
+            # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
+            # all experts into two fused 3-D tensors per layer:
+            #   experts.gate_up_proj  shape (E, 2*I, H)
+            #   experts.down_proj     shape (E, H, I)
+            # vLLM's Qwen3MoE load_weights still expects the older
+            # per-expert HF layout (experts.<i>.gate_proj.weight,
+            # experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
+            # so we un-fuse on the fly. Split order matches HF's forward:
+            #   gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
+            # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
+            params = self.model.state_dict()
+            for name in list(params.keys()):
+                param = params.pop(name)
+                if isinstance(param, DTensor):
+                    tensor = param.full_tensor().detach().contiguous()
+                else:
+                    tensor = param.detach().contiguous()
+                del param
+
+                if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
+                    prefix = name[: -len(".gate_up_proj")]
+                    num_experts, two_inter, _ = tensor.shape
+                    inter = two_inter // 2
+                    for i in range(num_experts):
+                        expert = tensor[i]
+                        yield (
+                            f"{prefix}.{i}.gate_proj.weight",
+                            expert[:inter].contiguous(),
+                        )
+                        yield (
+                            f"{prefix}.{i}.up_proj.weight",
+                            expert[inter:].contiguous(),
+                        )
+                    del tensor
+                elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
+                    prefix = name[: -len(".down_proj")]
+                    num_experts = tensor.shape[0]
+                    for i in range(num_experts):
+                        yield (
+                            f"{prefix}.{i}.down_proj.weight",
+                            tensor[i].contiguous(),
+                        )
+                    del tensor
+                else:
+                    yield name, tensor
+
+        trainer_args = IPCTrainerSendWeightsArgs(
+            send_mode="ray",
+            llm_handle=llm_handle,
+            packed=packed,
+            packed_buffer_size_bytes=1024 * 1024 * 1024,  # 1 GB
+        )
+        IPCWeightTransferEngine.trainer_send_weights(
+            iterator=_full_param_iter(),
+            trainer_args=trainer_args,
+        )
+
+
+@ray.remote(num_cpus=1)
+class DataParallelInferenceEngine:
+    """Manages a pool of DP-sharded vLLM LLM actors.
+
+    Spawns one MyLLM actor per placement group, waits for all engines to
+    finish initializing, and exposes generation / weight-sync helpers.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        pgs: list,
+        dp_master_ip: str,
+        dp_master_port: int,
+    ):
+        dp_size = len(pgs)
+        self.llm_actors = []
+        for r in range(dp_size):
+            sched = PlacementGroupSchedulingStrategy(
+                placement_group=pgs[r],
+                placement_group_capture_child_tasks=True,
+            )
+            actor = (
+                ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
+                .options(scheduling_strategy=sched)
+                .remote(
+                    model=model,
+                    enforce_eager=True,
+                    tensor_parallel_size=INFERENCE_TP_SIZE,
+                    distributed_executor_backend="ray",
+                    enable_expert_parallel=True,
+                    gpu_memory_utilization=0.35,
+                    weight_transfer_config=WeightTransferConfig(backend="ipc"),
+                    enable_sleep_mode=True,
+                    load_format="dummy",
+                    dp_rank=r,
+                    dp_size=dp_size,
+                    dp_master_ip=dp_master_ip,
+                    dp_master_port=dp_master_port,
+                )
+            )
+            self.llm_actors.append(actor)
+
+        ray.get([actor.ready.remote() for actor in self.llm_actors])
+
+    def get_llm_actors(self):
+        return self.llm_actors
+
+    def generate(self, prompts: list[str], sampling_params):
+        """Distribute prompts round-robin across DP ranks and collect results."""
+        dp_size = len(self.llm_actors)
+        per_rank: list[list[str]] = [[] for _ in range(dp_size)]
+        indices: list[list[int]] = [[] for _ in range(dp_size)]
+
+        for i, prompt in enumerate(prompts):
+            rank = i % dp_size
+            per_rank[rank].append(prompt)
+            indices[rank].append(i)
+
+        refs = [
+            actor.generate.remote(per_rank[r], sampling_params)
+            for r, actor in enumerate(self.llm_actors)
+            if per_rank[r]
+        ]
+        all_outputs = ray.get(refs)
+
+        ordered = [None] * len(prompts)
+        rank_idx = 0
+        for r in range(dp_size):
+            if per_rank[r]:
+                for local_i, orig_i in enumerate(indices[r]):
+                    ordered[orig_i] = all_outputs[rank_idx][local_i]
+                rank_idx += 1
+        return ordered
+
+    def init_weight_transfer(self):
+        ray.get(
+            [
+                actor.init_weight_transfer_engine.remote(
+                    dict(init_info=asdict(IPCWeightTransferInitInfo()))
+                )
+                for actor in self.llm_actors
+            ]
+        )
+
+    def start_weight_update(self, is_checkpoint_format: bool = True):
+        ray.get(
+            [
+                actor.start_weight_update.remote(
+                    is_checkpoint_format=is_checkpoint_format
+                )
+                for actor in self.llm_actors
+            ]
+        )
+
+    def finish_weight_update(self):
+        ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
+
+    def sleep(self, level: int = 0):
+        ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])
+
+    def wake_up(self, tags: list[str] | None = None):
+        ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])
+
+
+def main():
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
+            }
+        }
+    )
+
+    assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
+        "Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
+    )
+
+    local_model_path = snapshot_download(MODEL_NAME)
+    print(f"[init] Model downloaded to {local_model_path}")
+
+    fsdp_master_addr = get_ip()
+    fsdp_master_port = get_open_port()
+    dp_master_port = get_open_port()
+    dp_master_ip = get_ip()
+
+    # Create one placement group per DP rank (one GPU each).
+    pgs = []
+    for _ in range(INFERENCE_DP_SIZE):
+        pg = placement_group([{"GPU": 1, "CPU": 1}])
+        pgs.append(pg)
+    ray.get([pg.ready() for pg in pgs])
+    print(f"[init] {len(pgs)} placement groups ready.")
+
+    # Launch FSDP training workers, one per PG.
+    scheduling = [
+        PlacementGroupSchedulingStrategy(
+            placement_group=pgs[r],
+            placement_group_capture_child_tasks=True,
+        )
+        for r in range(FSDP_WORLD_SIZE)
+    ]
+
+    fsdp_workers = [
+        FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
+            local_model_path,
+            r,
+            FSDP_WORLD_SIZE,
+            fsdp_master_addr,
+            fsdp_master_port,
+        )
+        for r in range(FSDP_WORLD_SIZE)
+    ]
+    ray.get([w.get_rank.remote() for w in fsdp_workers])
+    print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")
+
+    # Launch DP inference engine (spawns and initializes all LLM actors).
+    inference_engine = DataParallelInferenceEngine.remote(
+        model=local_model_path,
+        pgs=pgs,
+        dp_master_ip=dp_master_ip,
+        dp_master_port=dp_master_port,
+    )
+    llm_actors = ray.get(inference_engine.get_llm_actors.remote())
+    print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")
+
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+    sampling_params = SamplingParams(temperature=0)
+
+    print("[generate] Generating with dummy weights...")
+    outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
+    print("-" * 60)
+    print("BEFORE weight sync (dummy weights):")
+    print("-" * 60)
+    for output in outputs:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+    # --- Weight transfer ---
+    print("[transfer] Initializing IPC weight transfer...")
+    ray.get(inference_engine.init_weight_transfer.remote())
+
+    # Two-phase sleep/wake pattern:
+    # 1. sleep(level=1) — offload weights to CPU, discard KV cache
+    # 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
+    # 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
+    # 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
+    print("[sync] Sleeping engines (offload weights + free KV cache)...")
+    ray.get(inference_engine.sleep.remote(level=1))
+
+    print("[sync] Waking weights (KV cache stays free)...")
+    ray.get(inference_engine.wake_up.remote(tags=["weights"]))
+
+    print("[sync] Starting weight update...")
+    ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
+
+    print("[sync] Packed IPC transfer FSDP → vLLM...")
+    ray.get(
+        [
+            w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
+            for w in fsdp_workers
+        ]
+    )
+
+    ray.get(inference_engine.finish_weight_update.remote())
+    print("[sync] Weight transfer complete.")
+
+    print("[sync] Waking KV cache + scheduling...")
+    ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))
+
+    print("[generate] Generating with synced weights...")
+    outputs_updated = ray.get(
+        inference_engine.generate.remote(prompts, sampling_params)
+    )
+    print("-" * 60)
+    print("AFTER weight sync (real weights):")
+    print("-" * 60)
+    for output in outputs_updated:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,384 @@
+{%- macro render_typescript_type(param_spec, required_params, is_nullable=false) -%}
+    {%- if param_spec.type == "array" -%}
+        {%- if param_spec['items'] -%}
+            {%- if param_spec['items']['type'] == "string" -%}
+                {{- "string[]" }}
+            {%- elif param_spec['items']['type'] == "number" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "integer" -%}
+                {{- "number[]" }}
+            {%- elif param_spec['items']['type'] == "boolean" -%}
+                {{- "boolean[]" }}
+            {%- else -%}
+                {%- set inner_type = render_typescript_type(param_spec['items'], required_params) -%}
+                {%- if inner_type == "object | object" or inner_type|length > 50 -%}
+                    {{- "any[]" }}
+                {%- else -%}
+                    {{- inner_type + "[]" }}
+                {%- endif -%}
+            {%- endif -%}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- else -%}
+            {{- "any[]" }}
+            {%- if param_spec.nullable -%}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type is defined and param_spec.type is iterable and param_spec.type is not string and param_spec.type is not mapping and param_spec.type[0] is defined -%}
+        {#- Handle array of types like ["object", "object"] from Union[dict, list] #}
+        {%- if param_spec.type | length > 1 -%}
+            {{- param_spec.type | join(" | ") }}
+        {%- else -%}
+            {{- param_spec.type[0] }}
+        {%- endif -%}
+    {%- elif param_spec.oneOf -%}
+        {#- Handle oneOf schemas - check for complex unions and fallback to any #}
+        {%- set has_object_variants = false -%}
+        {%- for variant in param_spec.oneOf -%}
+            {%- if variant.type == "object" -%}
+                {%- set has_object_variants = true -%}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- if has_object_variants and param_spec.oneOf|length > 1 -%}
+            {{- "any" }}
+        {%- else -%}
+            {%- for variant in param_spec.oneOf -%}
+                {{- render_typescript_type(variant, required_params) -}}
+                {%- if variant.description %}
+                    {{- "// " + variant.description }}
+                {%- endif -%}
+                {%- if variant.default is defined %}
+                    {{ "// default: " + variant.default|tojson }}
+                {%- endif -%}
+                {%- if not loop.last %}
+                    {{- " | " }}
+                {% endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "string" -%}
+        {%- if param_spec.enum -%}
+            {{- '"' + param_spec.enum|join('" | "') + '"' -}}
+        {%- else -%}
+            {{- "string" }}
+            {%- if param_spec.nullable %}
+                {{- " | null" }}
+            {%- endif -%}
+        {%- endif -%}
+    {%- elif param_spec.type == "number" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "integer" -%}
+        {{- "number" }}
+    {%- elif param_spec.type == "boolean" -%}
+        {{- "boolean" }}
+    {%- elif param_spec.type == "object" -%}
+        {%- if param_spec.properties -%}
+            {{- "{\n" }}
+            {%- for prop_name, prop_spec in param_spec.properties.items() -%}
+                {{- prop_name -}}
+                {%- if prop_name not in (param_spec.required or []) -%}
+                    {{- "?" }}
+                {%- endif -%}
+                {{- ": " }}
+                {{ render_typescript_type(prop_spec, param_spec.required or []) }}
+                {%- if not loop.last -%}
+                    {{-", " }}
+                {%- endif -%}
+            {%- endfor -%}
+            {{- "}" }}
+        {%- else -%}
+            {{- "object" }}
+        {%- endif -%}
+    {%- else -%}
+        {{- "any" }}
+    {%- endif -%}
+{%- endmacro -%}
+
+{%- macro render_tools(tools) -%}
+    {%- for tool in tools %}
+        {%- if tool.function is defined -%}
+            {#- Chat Completions format: {"type": "function", "function": {...}} #}
+            {%- set func = tool.function -%}
+            {%- if func.description is defined -%}
+                {{- "// " + func.description + "\n" }}
+            {%- endif -%}
+            {{- "type "+ func.name + " = " }}
+            {%- if func.parameters and func.parameters.properties %}
+                {{- "(_: {\n" }}
+                {%- for param_name, param_spec in func.parameters.properties.items() %}
+                    {%- if param_spec.description is defined %}
+                        {{- "// " + param_spec.description + "\n" }}
+                    {%- endif %}
+                    {{- param_name }}
+                    {%- if param_name not in (func.parameters.required or []) -%}
+                        {{- "?" }}
+                    {%- endif -%}
+                    {{- ": " }}
+                    {{- render_typescript_type(param_spec, func.parameters.required or []) }}
+                    {%- if param_spec.default is defined -%}
+                        {%- if param_spec.enum %}
+                            {{- ", // default: " + param_spec.default }}
+                        {%- elif param_spec.oneOf %}
+                            {{- "// default: " + param_spec.default }}
+                        {%- else %}
+                            {{- ", // default: " + param_spec.default|tojson }}
+                        {%- endif -%}
+                    {%- endif -%}
+                    {%- if not loop.last %}
+                        {{- ",\n" }}
+                    {%- else %}
+                        {{- "\n" }}
+                    {%- endif -%}
+                {%- endfor %}
+                {{- "}) => any;" }}
+            {%- else -%}
+                {{- "() => any;" }}
+            {%- endif -%}
+        {%- else -%}
+            {#- Responses format: {"type": "function", "name": "...", ...} #}
+            {%- if tool.description is defined -%}
+                {{- "// " + tool.description + "\n" }}
+            {%- endif -%}
+            {{- "type "+ tool.name + " = " }}
+            {%- if tool.parameters and tool.parameters.properties %}
+                {{- "(_: {\n" }}
+                {%- for param_name, param_spec in tool.parameters.properties.items() %}
+                    {%- if param_spec.description is defined %}
+                        {{- "// " + param_spec.description + "\n" }}
+                    {%- endif %}
+                    {{- param_name }}
+                    {%- if param_name not in (tool.parameters.required or []) -%}
+                        {{- "?" }}
+                    {%- endif -%}
+                    {{- ": " }}
+                    {{- render_typescript_type(param_spec, tool.parameters.required or []) }}
+                    {%- if param_spec.default is defined -%}
+                        {%- if param_spec.enum %}
+                            {{- ", // default: " + param_spec.default }}
+                        {%- elif param_spec.oneOf %}
+                            {{- "// default: " + param_spec.default }}
+                        {%- else %}
+                            {{- ", // default: " + param_spec.default|tojson }}
+                        {%- endif -%}
+                    {%- endif -%}
+                    {%- if not loop.last %}
+                        {{- ",\n" }}
+                    {%- else %}
+                        {{- "\n" }}
+                    {%- endif -%}
+                {%- endfor %}
+                {{- "}) => any;" }}
+            {%- else -%}
+                {{- "() => any;" }}
+            {%- endif -%}
+        {%- endif -%}
+        {%- if not loop.last -%}
+            {{- "\n" }}
+        {%- endif -%}
+    {%- endfor %}
+{%- endmacro -%}
+
+{{ bos_token }}
+
+{%- set system_token = '<|system_start|>' -%}
+{%- set end_system_token = '<|system_end|>' -%}
+{%- set developer_token = '<|developer_start|>' -%}
+{%- set end_developer_token = '<|developer_end|>' -%}
+{%- set user_token = '<|user_start|>' -%}
+{%- set end_user_token = '<|user_end|>' -%}
+{%- set assistant_token = '<|assistant_start|>' -%}
+{%- set end_assistant_token = '<|assistant_end|>' -%}
+{%- set inner_token = '<|inner_prefix|>' -%}
+{%- set outer_token = '<|inner_suffix|>' -%}
+{%- set tool_calls_token = '<|tools_prefix|>' -%}
+{%- set end_tool_calls_token = '<|tools_suffix|>' -%}
+
+{%- set ns = namespace(in_assistant=false, in_tool=false, in_inner=false, assistant_format=none) -%}
+
+{%- if messages and messages[0].role == 'system' -%}
+    {%- if "content" in messages[0] -%}
+        {%- if messages[0].content is string -%}
+            {{ system_token + messages[0].content + end_system_token }}
+        {%- elif messages[0].content is mapping and "text" in messages[0].content -%}
+            {{ system_token + messages[0].content.text + end_system_token }}
+        {%- else -%}
+            {{- raise_exception("Invalid system message") -}}
+        {%- endif -%}
+    {%- else -%}
+        {{- raise_exception("Invalid system message") -}}
+    {%- endif -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {{ system_token + 'You are Apertus, a helpful assistant created by the SwissAI initiative.\nKnowledge cutoff: 2024-04\nCurrent date: ' + strftime_now('%Y-%m-%d') + end_system_token }}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+
+{{ developer_token + 'Deliberation: ' }}
+{%- if enable_thinking is defined and enable_thinking -%}
+    {{ 'enabled\n' }}
+{%- else -%}
+    {{ 'disabled\n' }}
+{%- endif -%}
+{%- if tools is defined and tools -%}
+    {{ 'Tool Capabilities:\n' + render_tools(tools) }}
+{%- else -%}
+    {{ 'Tool Capabilities: disabled' }}
+{%- endif -%}
+{{ end_developer_token }}
+
+{%- for message in loop_messages -%}
+    {%- if message.role == 'user' -%}
+        {%- set ns.in_inner = false -%}
+        {%- if ns.in_tool -%}
+            {{ ']' }}
+            {%- set ns.in_tool = false -%}
+        {%- endif -%}
+        {%- if ns.in_assistant -%}
+            {{ end_assistant_token }}
+            {%- set ns.in_assistant = false -%}
+        {%- endif -%}
+        {%- if "content" in message -%}
+            {{ user_token }}
+            {%- if message.content is string -%}
+                {{ message.content }}
+            {%- elif message.content is mapping and "parts" in message.content -%}
+                {%- set parts = message.content.parts -%}
+                {%- for part in parts -%}
+                    {%- if part.type == "text" -%}
+                        {{ part.text }}
+                    {%- else -%}
+                        {{- raise_exception("Invalid user part: " + part.type) -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {{- raise_exception("Invalid user message: " + message.role) -}}
+            {%- endif -%}
+            {{ end_user_token }}
+        {%- endif -%}
+    {%- elif message.role == 'assistant' -%}
+        {%- if not ns.in_assistant -%}
+            {{ assistant_token }}
+            {%- set ns.in_assistant = true -%}
+        {%- endif -%}
+        {%- if "content" in message and message.content is not none -%}
+            {%- if message.content is string and (ns.assistant_format is none or ns.assistant_format == "string") -%}
+                {%- if ns.in_tool -%}
+                    {{ ']' }}
+                    {%- set ns.in_tool = false -%}
+                {%- endif -%}
+                {%- set ns.assistant_format = "string" -%}
+                {{ message.content }}
+            {%- elif message.content is mapping and "blocks" in message.content and (ns.assistant_format is none or ns.assistant_format == "mapping") -%}
+                {%- set ns.assistant_format = "mapping" -%}
+                {%- set blocks = message.content.blocks -%}
+                {%- for block in blocks -%}
+                    {%- if block.type == 'thoughts' -%}
+                        {%- if ns.in_tool -%}
+                            {{ ']' }}
+                            {%- set ns.in_tool = false -%}
+                        {%- endif -%}
+                        {%- if not ns.in_inner -%}
+                            {%- set ns.in_inner = true -%}
+                            {{ inner_token }}
+                        {%- endif -%}
+                        {{ block.text }}
+                    {%- elif block.type == 'tool_calls' -%}
+                        {%- if ns.in_tool -%}
+                            {{ ']' }}
+                            {%- set ns.in_tool = false -%}
+                        {%- endif -%}
+                        {%- if ns.in_inner and not loop.first and block.calls|length == 1 and block.calls[0].name == 'display_answers' -%}
+                            {%- set ns.in_inner = false -%}
+                            {{ outer_token }}
+                        {%- endif -%}
+                        {{ tool_calls_token + '[' }}
+                        {%- for tool_call in block.calls -%}
+                            {%- set args = tool_call.arguments -%}
+                            {%- if args is string -%}
+                                {{- '{"' + tool_call.name + '": ' + args + '}' }}
+                            {%- else -%}
+                                {{- '{"' + tool_call.name + '": ' + args|tojson + '}' }}
+                            {%- endif -%}
+                            {%- if not loop.last -%}
+                                {{- ", " }}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {{ ']' + end_tool_calls_token }}
+                    {%- elif block.type == 'tool_outputs' -%}
+                        {%- if ns.in_tool -%}
+                            {{- raise_exception("Cannot have both tool outputs as separate messages and tool outputs as blocks") -}}
+                        {%- endif -%}
+                        {{ '[' }}
+                        {%- for tool_output in block.outputs -%}
+                            {{- tool_output.output }}
+                            {%- if not loop.last -%}
+                                {{- ", " }}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {{- ']' }}
+                    {%- elif block.type == 'response' -%}
+                        {%- if ns.in_tool -%}
+                            {{ ']' }}
+                            {%- set ns.in_tool = false -%}
+                        {%- endif -%}
+                        {%- if (not loop.first and ns.in_inner) or (ns.in_assistant and ns.in_inner) -%}
+                            {%- set ns.in_inner = false -%}
+                            {{ outer_token }}
+                        {%- endif -%}
+                        {{ block.text }}
+                    {%- else -%}
+                        {{- raise_exception("Invalid assistant block type: " + block.type) -}}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- else -%}
+                {{- raise_exception("Invalid assistant content") -}}
+            {%- endif -%}
+        {%- elif not ("tool_calls" in message and message.tool_calls) -%}
+            {{- raise_exception("Invalid assistant message") -}}
+        {%- endif -%}
+        {%- if "tool_calls" in message and message.tool_calls -%}
+            {{ tool_calls_token + '[' }}
+            {%- for tool_call in message.tool_calls -%}
+                {%- if tool_call.type == 'function' -%}
+                    {%- set function = tool_call.function -%}
+                    {%- set args = function.arguments -%}
+                    {%- if args is string -%}
+                        {{- '{"' + function.name + '": ' + args + '}' }}
+                    {%- else -%}
+                        {{- '{"' + function.name + '": ' + args|tojson + '}' }}
+                    {%- endif -%}
+                    {%- if not loop.last -%}
+                        {{- ", " }}
+                    {%- endif -%}
+                {%- else -%}
+                    {{- raise_exception("Invalid tool call type: " + tool_call.type) -}}
+                {%- endif -%}
+            {%- endfor -%}
+            {{ ']' + end_tool_calls_token }}
+        {%- endif -%}
+    {%- elif message.role == 'tool' -%}
+        {%- if not ns.in_assistant -%}
+            {{- raise_exception("Tool message outside of assistant") -}}
+        {%- endif -%}
+        {%- if not ns.in_tool -%}
+            {{ '[' }}
+            {%- set ns.in_tool = true -%}
+        {%- else -%}
+            {{ ", "}}
+        {%- endif -%}
+        {%- if message.content is string -%}
+            {{ message.content }}
+        {%- else -%}
+            {{ message.content|tojson }}
+        {%- endif -%}
+    {%- else -%}
+        {{- raise_exception("Invalid message role") -}}
+    {%- endif -%}
+{%- endfor -%}
+{%- if ns.in_tool -%}
+    {{ ']' }}
+{%- endif -%}
+{%- if add_generation_prompt -%}
+    {{ assistant_token }}
+{%- endif -%}
@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.8.post1
-flashinfer-cubin==0.6.8.post1
+flashinfer-python==0.6.11.post2
+flashinfer-cubin==0.6.11.post2
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for packed tensor broadcasting functionality.

-Unit tests for packed_broadcast_producer and packed_broadcast_consumer.
+Unit tests for packed_nccl_broadcast_producer and packed_nccl_broadcast_consumer.
 These utilities enable efficient batched tensor transfer over NCCL.
 """

@@ -11,8 +11,12 @@ import torch

 from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
 from vllm.distributed.weight_transfer.packed_tensor import (
-    packed_broadcast_consumer,
-    packed_broadcast_producer,
+    pack_tensors,
+    packed_ipc_consumer,
+    packed_ipc_producer,
+    packed_nccl_broadcast_consumer,
+    packed_nccl_broadcast_producer,
+    unpack_tensor,
 )


@@ -90,91 +94,18 @@ class TestNCCLWeightTransferUpdateInfoPacked:
        assert info.packed is True


-# --- Unit Tests: packed_broadcast_producer ---
+# --- Unit Tests: packed_nccl_broadcast_producer ---


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 class TestPackedBroadcastProducer:
-    """Test packed_broadcast_producer function."""
-
-    def test_producer_broadcasts_tensors(self):
-        """Test that producer broadcasts all tensors."""
-        params = create_mock_model_params()
-        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
-
-        mock_group = MockCommunicationGroup()
-
-        # Use a small target size to force multiple batches
-        packed_broadcast_producer(
-            iterator=iter(params_cuda),
-            group=mock_group,
-            src=0,
-            post_iter_func=lambda x: x[1],
-            buffer_size_bytes=500,
-        )
-
-        # Should have broadcasted some tensors
-        assert mock_group.broadcast_count > 0
-        assert len(mock_group.broadcasted_tensors) > 0
-
-    def test_producer_single_large_tensor(self):
-        """Test with a single tensor larger than target size."""
-        # Create a large tensor
-        large_tensor = torch.randn(1000, 1000, dtype=torch.float32).cuda()
-        params = [("large_weight", large_tensor)]
-
-        mock_group = MockCommunicationGroup()
-
-        # Small target size to force the tensor to exceed it
-        packed_broadcast_producer(
-            iterator=iter(params),
-            group=mock_group,
-            src=0,
-            post_iter_func=lambda x: x[1],
-            buffer_size_bytes=100,
-        )
-
-        # Should still broadcast the tensor (at least 1 broadcast)
-        assert mock_group.broadcast_count >= 1
-        assert len(mock_group.broadcasted_tensors) >= 1
-
-        # Verify the total broadcasted size matches the tensor
-        expected_size = large_tensor.numel() * large_tensor.element_size()
-        actual_size = sum(t.numel() for t in mock_group.broadcasted_tensors)
-        assert actual_size == expected_size
-
-    def test_producer_multiple_batches(self):
-        """Test that tensors are properly batched when exceeding target size."""
-        # Create many small tensors
-        params = [
-            (f"weight_{i}", torch.randn(10, 10, dtype=torch.float32).cuda())
-            for i in range(20)
-        ]
-
-        mock_group = MockCommunicationGroup()
-
-        # Small target size to force multiple batches
-        packed_broadcast_producer(
-            iterator=iter(params),
-            group=mock_group,
-            src=0,
-            post_iter_func=lambda x: x[1],
-            buffer_size_bytes=2000,
-        )
-
-        # Should have multiple broadcasts
-        assert mock_group.broadcast_count > 1
-
-        # Total size should match sum of all tensors
-        expected_total = sum(t.numel() * t.element_size() for _, t in params)
-        actual_total = sum(t.numel() for t in mock_group.broadcasted_tensors)
-        assert actual_total == expected_total
+    """Test packed_nccl_broadcast_producer function."""

    def test_producer_empty_iterator(self):
        """Test producer handles empty iterator gracefully."""
        mock_group = MockCommunicationGroup()

-        packed_broadcast_producer(
+        packed_nccl_broadcast_producer(
            iterator=iter([]),
            group=mock_group,
            src=0,
@@ -186,64 +117,6 @@ class TestPackedBroadcastProducer:
        assert mock_group.broadcast_count == 0


-# --- Unit Tests: packed_broadcast_consumer ---
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-class TestPackedBroadcastConsumer:
-    """Test packed_broadcast_consumer function."""
-
-    def test_consumer_receives_tensors(self):
-        """Test that consumer receives and unpacks tensors."""
-        params = create_mock_model_params()
-        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
-
-        buffer_size = 2000
-
-        # First, run producer to get the broadcasted tensors
-        producer_group = MockCommunicationGroup()
-
-        packed_broadcast_producer(
-            iterator=iter(params_cuda),
-            group=producer_group,
-            src=0,
-            post_iter_func=lambda x: x[1],
-            buffer_size_bytes=buffer_size,
-        )
-
-        # Now run consumer with the broadcasted tensors
-        consumer_group = MockConsumerCommunicationGroup(
-            producer_group.broadcasted_tensors
-        )
-
-        state_dict_info = create_state_dict_info(params_cuda)
-
-        unpacked_tensors = {}
-
-        def post_unpack_func(tensor_list):
-            for name, tensor in tensor_list:
-                unpacked_tensors[name] = tensor.clone()
-
-        packed_broadcast_consumer(
-            iterator=iter(state_dict_info.items()),
-            group=consumer_group,
-            src=0,
-            post_unpack_func=post_unpack_func,
-            buffer_size_bytes=buffer_size,
-        )
-
-        # Verify all parameters were unpacked
-        assert len(unpacked_tensors) == len(params)
-
-        # Verify each tensor matches the original
-        for name, original_tensor in params_cuda:
-            assert name in unpacked_tensors
-            unpacked = unpacked_tensors[name]
-            assert unpacked.shape == original_tensor.shape
-            assert unpacked.dtype == original_tensor.dtype
-            assert torch.allclose(unpacked, original_tensor, rtol=1e-5, atol=1e-7)
-
-
 # --- Integration Tests: Producer-Consumer Roundtrip ---


@@ -260,7 +133,7 @@ class TestPackedBroadcastRoundtrip:
        buffer_size = 1000
        producer_group = MockCommunicationGroup()

-        packed_broadcast_producer(
+        packed_nccl_broadcast_producer(
            iterator=iter(params_cuda),
            group=producer_group,
            src=0,
@@ -279,7 +152,7 @@ class TestPackedBroadcastRoundtrip:
            for name, tensor in tensor_list:
                unpacked_tensors[name] = tensor.clone()

-        packed_broadcast_consumer(
+        packed_nccl_broadcast_consumer(
            iterator=iter(state_dict_info.items()),
            group=consumer_group,
            src=0,
@@ -306,7 +179,7 @@ class TestPackedBroadcastRoundtrip:
        buffer_size = 500
        producer_group = MockCommunicationGroup()

-        packed_broadcast_producer(
+        packed_nccl_broadcast_producer(
            iterator=iter(params),
            group=producer_group,
            src=0,
@@ -325,7 +198,7 @@ class TestPackedBroadcastRoundtrip:
            for name, tensor in tensor_list:
                unpacked_tensors[name] = tensor.clone()

-        packed_broadcast_consumer(
+        packed_nccl_broadcast_consumer(
            iterator=iter(state_dict_info.items()),
            group=consumer_group,
            src=0,
@@ -341,7 +214,7 @@ class TestPackedBroadcastRoundtrip:
            assert unpacked.dtype == original_tensor.dtype
            assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)

-    @pytest.mark.parametrize("target_size", [100, 1000, 10000, 100000])
+    @pytest.mark.parametrize("target_size", [100, 100000])
    def test_roundtrip_different_batch_sizes(self, target_size):
        """Test roundtrip with different target batch sizes."""
        params = create_mock_model_params(num_layers=5)
@@ -349,7 +222,7 @@ class TestPackedBroadcastRoundtrip:

        producer_group = MockCommunicationGroup()

-        packed_broadcast_producer(
+        packed_nccl_broadcast_producer(
            iterator=iter(params_cuda),
            group=producer_group,
            src=0,
@@ -368,7 +241,7 @@ class TestPackedBroadcastRoundtrip:
            for name, tensor in tensor_list:
                unpacked_tensors[name] = tensor.clone()

-        packed_broadcast_consumer(
+        packed_nccl_broadcast_consumer(
            iterator=iter(state_dict_info.items()),
            group=consumer_group,
            src=0,
@@ -407,7 +280,7 @@ class TestPackedBroadcastRoundtrip:
        buffer_size = 500
        producer_group = MockCommunicationGroup()

-        packed_broadcast_producer(
+        packed_nccl_broadcast_producer(
            iterator=iter(params),
            group=producer_group,
            src=0,
@@ -426,7 +299,7 @@ class TestPackedBroadcastRoundtrip:
            for name, tensor in tensor_list:
                unpacked_tensors[name] = tensor.clone()

-        packed_broadcast_consumer(
+        packed_nccl_broadcast_consumer(
            iterator=iter(state_dict_info.items()),
            group=consumer_group,
            src=0,
@@ -441,3 +314,462 @@ class TestPackedBroadcastRoundtrip:
            assert unpacked.shape == original_tensor.shape
            assert unpacked.dtype == original_tensor.dtype
            assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
+
+
+# --- Unit Tests: unpack_tensor ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestUnpackTensor:
+    """Test the shared unpack_tensor function."""
+
+    def test_unpack_produces_independent_copies(self):
+        """Verify unpacked tensors don't share memory with packed buffer."""
+        original = torch.randn(10, dtype=torch.float32).cuda()
+        packed = original.contiguous().view(torch.uint8).view(-1)
+
+        result = unpack_tensor(
+            packed,
+            names=["w"],
+            shapes=[[10]],
+            dtypes=[torch.float32],
+            tensor_sizes=[packed.numel()],
+        )
+
+        # Mutate the packed buffer
+        packed.zero_()
+
+        # Unpacked tensor should be unaffected
+        assert torch.allclose(result[0][1], original)
+
+
+# --- Unit Tests: pack_tensors ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackTensors:
+    """Test the shared pack_tensors function."""
+
+    def test_pack_basic(self):
+        """Test packing a few tensors into one buffer."""
+        params = [
+            ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
+            ("w2", torch.randn(5, dtype=torch.float16).cuda()),
+        ]
+
+        chunk = pack_tensors(
+            iterator=iter(params),
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=10_000_000,
+        )
+
+        assert chunk is not None
+        assert len(chunk.names) == 2
+        assert chunk.names == ["w1", "w2"]
+        assert chunk.shapes == [[10, 20], [5]]
+        assert chunk.dtypes == [torch.float32, torch.float16]
+        assert chunk.packed_tensor.dtype == torch.uint8
+
+    def test_pack_respects_buffer_limit(self):
+        """Test that packing stops when buffer_size_bytes is exceeded."""
+        params = [
+            (f"w{i}", torch.randn(100, 100, dtype=torch.float32).cuda())
+            for i in range(10)
+        ]
+
+        chunk = pack_tensors(
+            iterator=iter(params),
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=50_000,
+        )
+
+        assert chunk is not None
+        assert len(chunk.names) < 10
+
+    def test_pack_empty_iterator(self):
+        """Test that an empty iterator returns None."""
+        chunk = pack_tensors(
+            iterator=iter([]),
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=1000,
+        )
+        assert chunk is None
+
+    def test_pack_single_tensor_larger_than_buffer_warns(self):
+        """Test that a tensor exceeding buffer_size_bytes emits a warning."""
+        big = torch.randn(1000, 1000, dtype=torch.float32).cuda()
+        params = [("big", big)]
+
+        import warnings
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            chunk = pack_tensors(
+                iterator=iter(params),
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=100,
+            )
+        assert chunk is not None
+        assert len(chunk.names) == 1
+        assert any("exceeds buffer_size_bytes" in str(wi.message) for wi in w)
+
+    def test_pack_unpack_roundtrip(self):
+        """Test pack then unpack produces identical tensors."""
+        params = [
+            ("a", torch.randn(8, 16, dtype=torch.float32).cuda()),
+            ("b", torch.randn(4, dtype=torch.float16).cuda()),
+            ("c", torch.randn(3, 5, 7, dtype=torch.bfloat16).cuda()),
+        ]
+
+        chunk = pack_tensors(
+            iterator=iter(params),
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=10_000_000,
+        )
+
+        assert chunk is not None
+        result = unpack_tensor(
+            chunk.packed_tensor,
+            chunk.names,
+            chunk.shapes,
+            chunk.dtypes,
+            chunk.tensor_sizes,
+        )
+
+        assert len(result) == len(params)
+        for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
+            assert orig_name == res_name
+            assert res_tensor.shape == orig_tensor.shape
+            assert res_tensor.dtype == orig_tensor.dtype
+            assert torch.allclose(res_tensor, orig_tensor, rtol=1e-4, atol=1e-6)
+
+    def test_pack_multiple_chunks(self):
+        """Test consuming an iterator across multiple pack_tensors calls."""
+        params = [
+            (f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
+        ]
+        it = iter(params)
+
+        all_names = []
+        chunks = []
+        while True:
+            chunk = pack_tensors(it, lambda x: x[1], buffer_size_bytes=12_000)
+            if chunk is None:
+                break
+            chunks.append(chunk)
+            all_names.extend(chunk.names)
+
+        assert len(chunks) > 1
+        assert all_names == [f"w{i}" for i in range(6)]
+
+
+# --- Unit Tests: packed_ipc_producer ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackedIpcProducer:
+    """Test the packed_ipc_producer generator."""
+
+    def test_producer_yields_chunks(self):
+        """Test that the producer yields PackedIpcChunk objects."""
+        params = [
+            (f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
+        ]
+
+        chunks = list(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid="test-uuid",
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=12_000,
+            )
+        )
+
+        assert len(chunks) > 1
+
+    def test_producer_ipc_handle_has_uuid(self):
+        """Test that each chunk's ipc_handle is keyed by the given UUID."""
+        params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
+
+        chunks = list(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid="my-gpu-uuid",
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            )
+        )
+
+        assert "my-gpu-uuid" in chunks[0].ipc_handle
+
+    def test_producer_dtype_names_are_strings(self):
+        """Test that dtype_names are string representations."""
+        params = [
+            ("a", torch.randn(10, dtype=torch.float32).cuda()),
+            ("b", torch.randn(10, dtype=torch.float16).cuda()),
+        ]
+
+        chunks = list(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid="uuid",
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            )
+        )
+
+        assert chunks[0].dtype_names == ["float32", "float16"]
+
+    def test_producer_empty_iterator(self):
+        """Test producer with empty iterator yields nothing."""
+        chunks = list(
+            packed_ipc_producer(
+                iterator=iter([]),
+                gpu_uuid="uuid",
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=1000,
+            )
+        )
+        assert len(chunks) == 0
+
+
+# --- Integration Tests: IPC Producer-Consumer Roundtrip ---
+
+
+def _ipc_consumer_worker(cmd_q, ack_q, result_q, done_event, device_index):
+    """Worker that consumes chunks streamed one at a time from the parent.
+
+    CUDA IPC requires the consumer to be in a separate process from the
+    producer. The producer reuses a single IPC buffer between chunks, so
+    the parent must wait for our ack (sent after we copy the chunk to
+    CPU) before advancing the producer.
+    """
+    try:
+        torch.accelerator.set_device_index(device_index)
+        all_results = []
+        while True:
+            cd = cmd_q.get()
+            if cd is None:
+                break
+            result = packed_ipc_consumer(
+                ipc_handle=cd["ipc_handle"],
+                names=cd["names"],
+                shapes=cd["shapes"],
+                dtype_names=cd["dtype_names"],
+                tensor_sizes=cd["tensor_sizes"],
+                device_index=device_index,
+            )
+            # .cpu() forces a GPU→CPU copy off the shared IPC buffer, so
+            # the producer is free to overwrite it once we ack.
+            all_results.extend([(name, tensor.cpu()) for name, tensor in result])
+            del result
+            ack_q.put("ack")
+        result_q.put(("ok", all_results))
+    except Exception as e:
+        result_q.put(("error", str(e)))
+    # Keep the process alive until the parent has finished reading from
+    # the result queue — torch serializes CPU tensors via fd sharing,
+    # which requires this process's resource-sharer server to be running.
+    done_event.wait(timeout=60)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackedIpcRoundtrip:
+    """Test IPC producer-consumer roundtrip using real CUDA IPC.
+
+    These tests spawn a child process for the consumer because
+    rebuild_cuda_tensor requires a separate process from the one that
+    called reduce_tensor.
+    """
+
+    def _get_gpu_uuid(self) -> str:
+        device_index = torch.cuda.current_device()
+        props = torch.cuda.get_device_properties(device_index)
+        return str(props.uuid)
+
+    def _run_roundtrip(self, chunk_iter, device_index, timeout=30):
+        """Stream chunks through a child consumer one at a time.
+
+        ``packed_ipc_producer`` reuses a single IPC buffer for every
+        chunk, so the producer must not be advanced until the consumer
+        has finished reading the current chunk. We enforce that with an
+        ack queue: the consumer puts ``"ack"`` after it has copied the
+        chunk to CPU, and only then do we pull the next chunk from the
+        generator.
+
+        Returns ``(num_chunks, results)``.
+        """
+        import multiprocessing as mp
+
+        ctx = mp.get_context("spawn")
+        cmd_q = ctx.Queue()
+        ack_q = ctx.Queue()
+        result_q = ctx.Queue()
+        done_event = ctx.Event()
+        proc = ctx.Process(
+            target=_ipc_consumer_worker,
+            args=(cmd_q, ack_q, result_q, done_event, device_index),
+        )
+        proc.start()
+
+        num_chunks = 0
+        try:
+            for chunk in chunk_iter:
+                cmd_q.put(
+                    {
+                        "ipc_handle": chunk.ipc_handle,
+                        "names": chunk.names,
+                        "shapes": chunk.shapes,
+                        "dtype_names": chunk.dtype_names,
+                        "tensor_sizes": chunk.tensor_sizes,
+                    }
+                )
+                if ack_q.get(timeout=timeout) != "ack":
+                    raise RuntimeError("Consumer did not ack chunk")
+                num_chunks += 1
+            cmd_q.put(None)
+            status, payload = result_q.get(timeout=timeout)
+        finally:
+            done_event.set()
+            proc.join(timeout=10)
+            if proc.is_alive():
+                proc.kill()
+
+        if status == "error":
+            raise RuntimeError(f"Consumer process failed: {payload}")
+        # Reclaim IPC-shared memory now that the child has released it
+        torch.cuda.ipc_collect()
+        return num_chunks, payload
+
+    def test_roundtrip_basic(self):
+        """Test basic IPC producer -> consumer roundtrip."""
+        params = [
+            ("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
+            ("w2", torch.randn(5, dtype=torch.float16).cuda()),
+        ]
+        gpu_uuid = self._get_gpu_uuid()
+        device_index = torch.cuda.current_device()
+
+        num_chunks, result = self._run_roundtrip(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid=gpu_uuid,
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            ),
+            device_index,
+        )
+
+        assert num_chunks == 1
+        assert len(result) == 2
+        for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
+            assert orig_name == res_name
+            assert res_tensor.shape == orig_tensor.shape
+            assert res_tensor.dtype == orig_tensor.dtype
+            assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_roundtrip_dtypes(self, dtype):
+        """Test IPC roundtrip with different dtypes."""
+        params = create_mock_model_params(num_layers=2, dtype=dtype)
+        params_cuda = [(n, t.cuda()) for n, t in params]
+        gpu_uuid = self._get_gpu_uuid()
+        device_index = torch.cuda.current_device()
+
+        _, result = self._run_roundtrip(
+            packed_ipc_producer(
+                iterator=iter(params_cuda),
+                gpu_uuid=gpu_uuid,
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            ),
+            device_index,
+        )
+
+        assert len(result) == len(params_cuda)
+        for (orig_name, orig_tensor), (res_name, res_tensor) in zip(
+            params_cuda, result
+        ):
+            assert orig_name == res_name
+            assert res_tensor.dtype == dtype
+            assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
+
+    def test_roundtrip_multiple_chunks(self):
+        """Test IPC roundtrip across multiple chunks."""
+        params = [
+            (f"layer{i}.weight", torch.randn(100, 100, dtype=torch.float32).cuda())
+            for i in range(8)
+        ]
+        gpu_uuid = self._get_gpu_uuid()
+        device_index = torch.cuda.current_device()
+
+        num_chunks, result = self._run_roundtrip(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid=gpu_uuid,
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=50_000,
+            ),
+            device_index,
+        )
+
+        assert num_chunks > 1
+        assert len(result) == len(params)
+        for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
+            assert orig_name == res_name
+            assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-5, atol=1e-7)
+
+    def test_roundtrip_non_contiguous(self):
+        """Test IPC roundtrip with non-contiguous tensors."""
+        params = [
+            ("transposed", torch.randn(20, 10, dtype=torch.float32).cuda().T),
+            ("sliced", torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]),
+        ]
+        gpu_uuid = self._get_gpu_uuid()
+        device_index = torch.cuda.current_device()
+
+        for _, t in params:
+            assert not t.is_contiguous()
+
+        _, result = self._run_roundtrip(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid=gpu_uuid,
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            ),
+            device_index,
+        )
+
+        for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
+            assert orig_name == res_name
+            assert res_tensor.shape == orig_tensor.shape
+            assert res_tensor.dtype == orig_tensor.dtype
+            assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
+
+    def test_consumer_wrong_uuid_raises(self):
+        """Test that consumer raises ValueError for unknown GPU UUID."""
+        params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
+        gpu_uuid = self._get_gpu_uuid()
+
+        chunks = list(
+            packed_ipc_producer(
+                iterator=iter(params),
+                gpu_uuid=gpu_uuid,
+                post_iter_func=lambda x: x[1],
+                buffer_size_bytes=10_000_000,
+            )
+        )
+
+        c = chunks[0]
+        fake_handle = {"fake-uuid-12345": c.ipc_handle[gpu_uuid]}
+
+        with pytest.raises(ValueError, match="IPC handle not found"):
+            packed_ipc_consumer(
+                ipc_handle=fake_handle,
+                names=c.names,
+                shapes=c.shapes,
+                dtype_names=c.dtype_names,
+                tensor_sizes=c.tensor_sizes,
+                device_index=torch.cuda.current_device(),
+            )
@@ -389,7 +389,7 @@ class TestIPCWeightTransferUpdateInfoValidation:

        # Create a dummy tensor and IPC handle
        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
+        _, ipc_handle = reduce_tensor(dummy_tensor)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
        ipc_handles = [{gpu_uuid: ipc_handle}]

@@ -410,7 +410,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
            pytest.skip("Need at least 1 GPU for this test")

        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
+        _, ipc_handle = reduce_tensor(dummy_tensor)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
        ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]

@@ -428,7 +428,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
            pytest.skip("Need at least 1 GPU for this test")

        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
+        _, ipc_handle = reduce_tensor(dummy_tensor)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
        ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]

@@ -446,7 +446,7 @@ class TestIPCWeightTransferUpdateInfoValidation:
            pytest.skip("Need at least 1 GPU for this test")

        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
+        _, ipc_handle = reduce_tensor(dummy_tensor)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
        ipc_handles = [{gpu_uuid: ipc_handle}]  # Only one handle

@@ -458,65 +458,9 @@ class TestIPCWeightTransferUpdateInfoValidation:
                ipc_handles=ipc_handles,
            )

-    def test_valid_update_info_from_pickled(self, monkeypatch):
-        """Test creating IPCWeightTransferUpdateInfo from pickled handles."""
-        if torch.accelerator.device_count() < 1:
-            pytest.skip("Need at least 1 GPU for this test")
-
-        monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
-
-        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
-        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
-        ipc_handles = [{gpu_uuid: ipc_handle}]
-
-        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
-
-        info = IPCWeightTransferUpdateInfo(
-            names=["layer.weight"],
-            dtype_names=["float32"],
-            shapes=[[10, 10]],
-            ipc_handles_pickled=pickled,
-        )
-        assert info.ipc_handles == ipc_handles
-        assert info.ipc_handles_pickled is None
-
-    def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
-        """Test that pickled handles are rejected unless env flag is enabled."""
-        monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
-
-        with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
-            IPCWeightTransferUpdateInfo(
-                names=[],
-                dtype_names=[],
-                shapes=[],
-                ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
-            )
-
-    def test_both_handles_and_pickled_raises(self):
-        """Test that providing both ipc_handles and ipc_handles_pickled raises."""
-        if torch.accelerator.device_count() < 1:
-            pytest.skip("Need at least 1 GPU for this test")
-
-        dummy_tensor = torch.ones(10, 10, device="cuda:0")
-        ipc_handle = reduce_tensor(dummy_tensor)
-        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
-        ipc_handles = [{gpu_uuid: ipc_handle}]
-
-        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
-
-        with pytest.raises(ValueError, match="Cannot specify both"):
-            IPCWeightTransferUpdateInfo(
-                names=["layer.weight"],
-                dtype_names=["float32"],
-                shapes=[[10, 10]],
-                ipc_handles=ipc_handles,
-                ipc_handles_pickled=pickled,
-            )
-
-    def test_neither_handles_nor_pickled_raises(self):
-        """Test that providing neither ipc_handles nor ipc_handles_pickled raises."""
-        with pytest.raises(ValueError, match="must be provided"):
+    def test_missing_ipc_handles_raises(self):
+        """Test that omitting ipc_handles raises TypeError."""
+        with pytest.raises(TypeError):
            IPCWeightTransferUpdateInfo(
                names=["layer.weight"],
                dtype_names=["float32"],
@@ -552,10 +496,10 @@ class TestIPCEngineParsing:
        # Create dummy IPC handles
        dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
        dummy_tensor2 = torch.ones(50, device="cuda:0")
-        ipc_handle1 = reduce_tensor(dummy_tensor1)
-        ipc_handle2 = reduce_tensor(dummy_tensor2)
+        _, ipc_args1 = reduce_tensor(dummy_tensor1)
+        _, ipc_args2 = reduce_tensor(dummy_tensor2)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
-        ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
+        ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]

        update_info = engine.parse_update_info(
            {
@@ -585,10 +529,10 @@ class TestIPCEngineParsing:

        dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
        dummy_tensor2 = torch.ones(50, device="cuda:0")
-        ipc_handle1 = reduce_tensor(dummy_tensor1)
-        ipc_handle2 = reduce_tensor(dummy_tensor2)
+        _, ipc_args1 = reduce_tensor(dummy_tensor1)
+        _, ipc_args2 = reduce_tensor(dummy_tensor2)
        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
-        ipc_handles = [{gpu_uuid: ipc_handle1}, {gpu_uuid: ipc_handle2}]
+        ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]

        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")

@@ -604,10 +548,36 @@ class TestIPCEngineParsing:
        assert isinstance(update_info, IPCWeightTransferUpdateInfo)
        assert update_info.names == ["w1", "w2"]
        assert len(update_info.ipc_handles) == 2
-        assert update_info.ipc_handles_pickled is None
        assert gpu_uuid in update_info.ipc_handles[0]
        assert gpu_uuid in update_info.ipc_handles[1]

+    def test_parse_update_info_both_handles_and_pickled_raises(self):
+        """Test that providing both ipc_handles and ipc_handles_pickled raises."""
+        if torch.accelerator.device_count() < 1:
+            pytest.skip("Need at least 1 GPU for this test")
+
+        config = WeightTransferConfig(backend="ipc")
+        parallel_config = create_mock_parallel_config()
+        engine = IPCWeightTransferEngine(config, parallel_config)
+
+        dummy_tensor = torch.ones(10, 10, device="cuda:0")
+        _, ipc_handle = reduce_tensor(dummy_tensor)
+        gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
+        ipc_handles = [{gpu_uuid: ipc_handle}]
+
+        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
+
+        with pytest.raises(ValueError, match="Cannot specify both"):
+            engine.parse_update_info(
+                {
+                    "names": ["layer.weight"],
+                    "dtype_names": ["float32"],
+                    "shapes": [[10, 10]],
+                    "ipc_handles": ipc_handles,
+                    "ipc_handles_pickled": pickled,
+                }
+            )
+

 # --- Integration Test: IPC Weight Transfer Between Ray Tasks ---

@@ -629,13 +599,15 @@ class TrainerActor:
        self.tensor.fill_(42.0)  # Fill with 42 to verify correct transfer

        # Create IPC handle (tensor must stay alive for IPC to work)
-        ipc_handle = reduce_tensor(self.tensor)
+        # reduce_tensor returns (rebuild_func, args); we only send args
+        # since the receiver imports rebuild_cuda_tensor directly.
+        _, ipc_args = reduce_tensor(self.tensor)
        gpu_uuid = get_physical_gpu_id(0)

        torch.accelerator.synchronize()

        self.ipc_handle_dict = {
-            "ipc_handle": ipc_handle,
+            "ipc_handle": ipc_args,
            "gpu_uuid": gpu_uuid,
            "shape": tensor_shape,
            "dtype": tensor_dtype,
@@ -652,6 +624,12 @@ def inference_receive_ipc_tensor(
    mode: str = "ray",
 ) -> dict:
    """Inference task that receives tensor via IPCWeightTransferEngine."""
+    import os
+
+    # Worker-side: ipc_handles_pickled is deserialized via pickle.
+    if mode == "http":
+        os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
    from unittest.mock import MagicMock

    import torch
@@ -684,7 +662,6 @@ def inference_receive_ipc_tensor(
            # Clone tensor to keep it after engine cleans up
            received_tensors.append((name, tensor.clone()))

-    # Build update dict and go through parse_update_info (exercises __post_init__)
    ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]

    if mode == "ray":
@@ -695,6 +672,7 @@ def inference_receive_ipc_tensor(
            "ipc_handles": ipc_handles,
        }
    elif mode == "http":
+        # Simulate HTTP transport: pickle + base64 encode handles
        pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
        update_dict = {
            "names": ["test.weight"],
@@ -743,7 +721,8 @@ def test_ipc_weight_transfer_between_processes(mode: str):

    Parametrized over transport modes:
    - 'ray':  ipc_handles passed directly.
-    - 'http': ipc_handles pickled + base64-encoded, unpickled via __post_init__.
+    - 'http': ipc_handles pickled + base64-encoded, deserialized in
+              parse_update_info before constructing the dataclass.

    IPC requires same-GPU access, so we use a placement group to co-locate
    the trainer actor and inference task on the same GPU.
@@ -801,7 +780,7 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises():

    # Create IPC handle with wrong GPU UUID
    dummy_tensor = torch.ones(10, 10, device="cuda:0")
-    ipc_handle = reduce_tensor(dummy_tensor)
+    _, ipc_handle = reduce_tensor(dummy_tensor)
    wrong_uuid = "wrong-uuid-12345"
    ipc_handles = [{wrong_uuid: ipc_handle}]

@@ -86,13 +86,56 @@ async def test_passed_api_token(server: RemoteOpenAIServer):
    indirect=True,
 )
@pytest.mark.asyncio
-async def test_not_v1_api_token(server: RemoteOpenAIServer):
-    # Authorization check is skipped for any paths that
-    # don't start with /v1 (e.g. /v1/chat/completions).
+async def test_not_v1_or_v2_path_skips_auth(server: RemoteOpenAIServer):
+    # Authorization check is skipped for paths that
+    # don't start with /v1 or /v2 (e.g. /health, /metrics).
    response = requests.get(server.url_for("health"))
    assert response.status_code == HTTPStatus.OK


+# ---------------------------------------------------------------------------
+# /v2 path authentication tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "server",
+    [["--api-key", "test"]],
+    indirect=True,
+)
+@pytest.mark.asyncio
+async def test_v2_endpoint_rejects_missing_api_token(server: RemoteOpenAIServer):
+    # /v2/embed should require authentication when --api-key is set.
+    body = {
+        "model": MODEL_NAME,
+        "texts": ["hello"],
+        "embedding_types": ["float"],
+    }
+    response = requests.post(server.url_for("/v2/embed"), json=body)
+    assert response.status_code == HTTPStatus.UNAUTHORIZED
+
+
+@pytest.mark.parametrize(
+    "server",
+    [["--api-key", "test"]],
+    indirect=True,
+)
+@pytest.mark.asyncio
+async def test_v2_endpoint_accepts_valid_api_token(server: RemoteOpenAIServer):
+    # /v2/embed should accept requests with a valid API key.
+    body = {
+        "model": MODEL_NAME,
+        "texts": ["hello"],
+        "embedding_types": ["float"],
+    }
+    response = requests.post(
+        server.url_for("/v2/embed"),
+        json=body,
+        headers={"Authorization": "Bearer test"},
+    )
+    assert response.status_code == HTTPStatus.OK
+
+
@pytest.mark.parametrize(
    "server",
    ["--enable-request-id-headers"],
@@ -0,0 +1,12 @@
+model_name: "amd/Qwen3.5-35B-A3B-MXFP4"
+accuracy_threshold: 0.89
+tolerance: 0.03
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --max-model-len 4096
+  --tensor-parallel-size 2
+  --gpu-memory-utilization 0.35
+  --moe-backend aiter
+env:
+  VLLM_ROCM_USE_AITER: "1"
@@ -1,8 +1,10 @@
 model_name: "amd/Qwen3.5-35B-A3B-MXFP4"
-accuracy_threshold: 0.82
+accuracy_threshold: 0.89
 tolerance: 0.03
 num_questions: 1319
 num_fewshot: 5
 server_args: >-
  --max-model-len 4096
  --tensor-parallel-size 2
+  --moe-backend emulation
+  --gpu-memory-utilization 0.35
@@ -3,4 +3,5 @@ DeepSeek-R1-DP_MI325.yaml
 DeepSeek-V3.2-TP_MI325.yaml
 DeepSeek-V3.2-DP_MI325.yaml
 Qwen3-30B-A3B-NVFP4.yaml
-Qwen3.5-35B-A3B-MXFP4-TP2.yaml
+Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
+Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
@@ -1,2 +1,3 @@
 Qwen3.5-35B-A3B-DEP2.yaml
-Qwen3.5-35B-A3B-MXFP4-TP2.yaml
+Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
+Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
@@ -272,6 +272,31 @@ class TestIrOpImplDispatch:
        # Restored to empty
        assert _custom_add.get_priority() == []

+    @pytest.mark.parametrize(
+        "default,override",
+        [
+            (["impl_even", "impl_b"], ["impl_a"]),
+            (["impl_a"], ["impl_even", "impl_b"]),
+        ],
+    )
+    def test_set_default_priority(
+        self, custom_add_op, default: list[str], override: list[str]
+    ):
+        _custom_add = custom_add_op
+        assert _custom_add.get_priority() == []
+
+        _custom_add.set_default(default)
+        assert _custom_add.get_priority() == default
+
+        # Priority doesn't change after exiting the set_priority context.
+        with _custom_add.set_priority(override):
+            assert _custom_add.get_priority() == override
+        assert _custom_add.get_priority() == default
+
+        # Should override the previous default.
+        _custom_add.set_default(override)
+        assert _custom_add.get_priority() == override
+
    @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
    def test_dispatch_priority_order(self, custom_add_op, overload: str):
        _custom_add = custom_add_op
@@ -381,6 +406,26 @@ class TestIrOpImplDispatch:
        assert "priority not set" in message


+@pytest.mark.parametrize("default", [True, False])
+def test_set_default_torch_wrap(default: bool):
+    """set_default_torch_wrap permanently flips the global flag."""
+    original = vllm.ir.op._ENABLE_TORCH_WRAP
+    try:
+        vllm.ir.set_default_torch_wrap(default)
+        assert vllm.ir.op._ENABLE_TORCH_WRAP is default
+
+        # Flag doesn't change after exiting the enable_torch_wrap context.
+        with vllm.ir.enable_torch_wrap(not default):
+            assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
+        assert vllm.ir.op._ENABLE_TORCH_WRAP is default
+
+        # Should override the previous default.
+        vllm.ir.set_default_torch_wrap(not default)
+        assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
+    finally:
+        vllm.ir.op._ENABLE_TORCH_WRAP = original
+
+
@pytest.fixture
 def custom_mm_op(fake_vllm_ir):
    """Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""
@@ -514,3 +514,35 @@ def test_non_causal_autoselect_backend():
            kv_cache_dtype=None,
        )
        assert backend.supports_non_causal()
+
+
+@pytest.mark.parametrize(
+    "kv_cache_dtype",
+    [
+        "fp8_e5m2",
+        "fp8_ds_mla",
+        "fp8_inc",
+        "nvfp4",
+        "fp8_per_token_head",
+        "int8_per_token_head",
+    ],
+)
+def test_flash_attn_rejects_unhandled_kv_cache_dtypes(kv_cache_dtype: str):
+    """FlashAttentionBackend must not claim support for kv_cache dtypes
+    that it cannot handle."""
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
+    assert not FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["fp8", "fp8_e4m3"])
+def test_flash_attn_accepts_handled_fp8_variants(
+    kv_cache_dtype: str, monkeypatch: pytest.MonkeyPatch
+):
+    """FlashAttentionBackend must accept the two fp8 dtypes it can actually
+    handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
+    import vllm.v1.attention.backends.flash_attn as fa_mod
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
+    monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
+    assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
@@ -1,256 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for CPU FP8 W8A16 block-scaled fused MoE kernel."""
-
-import math
-import sys
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm.platforms import current_platform
-from vllm.utils.torch_utils import set_random_seed
-
-if not current_platform.is_cpu():
-    pytest.skip("skipping CPU-only tests", allow_module_level=True)
-
-import vllm._custom_ops as ops  # noqa: E402
-
-if not hasattr(torch.ops._C, "fused_experts_cpu"):
-    pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
-
-
-BLOCK_SIZE = [128, 128]  # [block_n, block_k]
-
-# FP8 weight generation parameters
-_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
-FP8_SCALE = _FP8_INFO.max  # 448.0
-FACTOR_FOR_SCALE = 1e-3
-
-# Tolerance for FP8 W8A16
-FP8_W8A16_ATOL = 1e-2
-FP8_W8A16_RTOL = 1e-2
-
-
-def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
-    d = x.shape[-1] // 2
-    return F.silu(x[..., :d]) * x[..., d:]
-
-
-def _block_dequant_weight(
-    weight: torch.Tensor,
-    scales: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-dequantize FP8 weight [E, N, K] → float [E, N, K].
-
-    Each (block_n × block_k) tile is multiplied by its per-block scale.
-    """
-    E, N, K = weight.shape
-    block_n, block_k = block_size
-    pad_N = (block_n - N % block_n) % block_n
-    pad_K = (block_k - K % block_k) % block_k
-
-    if pad_N > 0 or pad_K > 0:
-        weight = F.pad(weight, (0, pad_K, 0, pad_N))
-
-    n_tiles = math.ceil(N / block_n)
-    k_tiles = math.ceil(K / block_k)
-
-    weight_block = (
-        weight.view(E, n_tiles, block_n, k_tiles, block_k)
-        .permute(0, 1, 3, 2, 4)
-        .float()
-        .contiguous()
-    )
-    weight_scaled = (
-        (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
-        .permute(0, 1, 3, 2, 4)
-        .contiguous()
-    )
-    if pad_N > 0 or pad_K > 0:
-        weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
-        weight_scaled = weight_scaled[..., :N, :K].contiguous()
-    else:
-        weight_scaled = weight_scaled.view(E, N, K)
-    return weight_scaled
-
-
-def ref_w8a16_block_fp8_moe(
-    a: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    w1_s: torch.Tensor,
-    w2_s: torch.Tensor,
-    topk_weight: torch.Tensor,
-    topk_ids: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Reference FP8 W8A16 block-scaled fused MoE in pure torch.
-
-    Steps:
-      1. Block-dequant FP8 weights → float
-      2. For each expert: matmul → SiLU+Mul → matmul
-      3. Weighted sum across top-k experts
-    """
-    B, D = a.shape
-    topk = topk_ids.size(1)
-
-    w1_dq = _block_dequant_weight(w1, w1_s, block_size)
-    w2_dq = _block_dequant_weight(w2, w2_s, block_size)
-
-    a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
-    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
-
-    topk_weight_flat = topk_weight.view(-1)
-    topk_ids_flat = topk_ids.view(-1)
-
-    for i in range(w1_dq.shape[0]):
-        mask = topk_ids_flat == i
-        if mask.sum():
-            ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
-            ic1 = _silu_and_mul(ic0)
-            out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
-
-    return (
-        (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
-        .sum(dim=1)
-        .to(a.dtype)
-    )
-
-
-def _make_fp8_moe_weights(
-    E: int,
-    N: int,
-    K: int,
-    block_size: list[int],
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Generate random FP8 MoE weights with random block scales.
-
-    Weight generation follows SGLang: ``randn * FP8_SCALE`` → clamp → cast.
-    Scales are small random values (``FACTOR_FOR_SCALE``), independent of
-    the actual weight magnitudes — this is sufficient to test the kernel's
-    block-dequant + matmul correctness.
-
-    Returns: (w1, w2, w1_s, w2_s)
-    """
-    block_n, block_k = block_size
-
-    w1 = (
-        (torch.randn(E, 2 * N, K) * FP8_SCALE)
-        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
-        .to(torch.float8_e4m3fn)
-    )
-    w2 = (
-        (torch.randn(E, K, N) * FP8_SCALE)
-        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
-        .to(torch.float8_e4m3fn)
-    )
-
-    w1_s = (
-        torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
-        * FACTOR_FOR_SCALE
-    )
-    w2_s = (
-        torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
-        * FACTOR_FOR_SCALE
-    )
-    return w1, w2, w1_s, w2_s
-
-
-def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
-    """VNNI-prepack each expert's weight via ``convert_weight_packed``."""
-    return torch.stack(
-        [torch.ops._C.convert_weight_packed(w[e]) for e in range(w.shape[0])]
-    )
-
-
-NUM_TOKENS = [1, 2, 64, 121]
-# (M, intermediate_size N, hidden_size K, num_experts E, topk)
-MoE_CONFIGS = [
-    (256, 512, 8, 2),
-    (256, 512, 8, 4),
-    (512, 256, 8, 2),
-    (512, 256, 8, 4),
-    (512, 512, 8, 2),
-    (512, 512, 8, 4),
-    (768, 2048, 8, 2),
-    (768, 2048, 8, 4),
-    (768, 2048, 128, 8),
-]
-SEEDS = [0]
-
-
-@pytest.mark.parametrize("M", NUM_TOKENS)
-@pytest.mark.parametrize("N,K,E,topk", MoE_CONFIGS)
-@pytest.mark.parametrize("seed", SEEDS)
-def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
-    """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
-    set_random_seed(seed)
-
-    a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
-    w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
-
-    score = torch.randn(M, E, dtype=torch.bfloat16)
-    score = torch.softmax(score, dim=-1, dtype=torch.float32)
-    topk_weight, topk_ids = torch.topk(score, topk)
-    topk_ids = topk_ids.to(torch.int32)
-
-    ref_out = ref_w8a16_block_fp8_moe(
-        a,
-        w1,
-        w2,
-        w1_s,
-        w2_s,
-        topk_weight,
-        topk_ids,
-        BLOCK_SIZE,
-    )
-
-    pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
-
-    # Test inplace=False against reference
-    out = ops.fused_experts_cpu(
-        a.clone(),
-        pw1,
-        pw2,
-        topk_weight,
-        topk_ids,
-        False,
-        ops.CPUQuantMethod.FP8_W8A16,
-        w1_s,
-        w2_s,
-        None,
-        None,
-        BLOCK_SIZE,
-        True,
-    )
-    torch.testing.assert_close(
-        ref_out.bfloat16(),
-        out,
-        atol=FP8_W8A16_ATOL,
-        rtol=FP8_W8A16_RTOL,
-    )
-
-    # Test inplace=True produces identical output
-    out_inplace = ops.fused_experts_cpu(
-        a.clone(),
-        pw1,
-        pw2,
-        topk_weight,
-        topk_ids,
-        True,
-        ops.CPUQuantMethod.FP8_W8A16,
-        w1_s,
-        w2_s,
-        None,
-        None,
-        BLOCK_SIZE,
-        True,
-    )
-    torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))
@@ -0,0 +1,500 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for CPU quantized fused MoE kernels (FP8 W8A16 and MXFP4 W4A16)."""
+
+import math
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+if not current_platform.is_cpu():
+    pytest.skip("skipping CPU-only tests", allow_module_level=True)
+
+import vllm._custom_ops as ops  # noqa: E402
+
+if not hasattr(torch.ops._C, "fused_experts_cpu"):
+    pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
+
+
+def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
+    """VNNI-prepack expert weights via ``convert_weight_packed``."""
+    return torch.ops._C.convert_weight_packed(w)
+
+
+# FP8 W8A16 block-scaled fused MoE
+
+BLOCK_SIZE = [128, 128]  # [block_n, block_k]
+
+_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
+FP8_SCALE = _FP8_INFO.max  # 448.0
+FACTOR_FOR_SCALE = 1e-3
+
+
+def _block_dequant_weight(
+    weight: torch.Tensor,
+    scales: torch.Tensor,
+    block_size: list[int],
+) -> torch.Tensor:
+    """Block-dequantize FP8 weight [E, N, K] -> float [E, N, K]."""
+    E, N, K = weight.shape
+    block_n, block_k = block_size
+    pad_N = (block_n - N % block_n) % block_n
+    pad_K = (block_k - K % block_k) % block_k
+
+    if pad_N > 0 or pad_K > 0:
+        weight = F.pad(weight, (0, pad_K, 0, pad_N))
+
+    n_tiles = math.ceil(N / block_n)
+    k_tiles = math.ceil(K / block_k)
+
+    weight_block = (
+        weight.view(E, n_tiles, block_n, k_tiles, block_k)
+        .permute(0, 1, 3, 2, 4)
+        .float()
+        .contiguous()
+    )
+    weight_scaled = (
+        (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
+        .permute(0, 1, 3, 2, 4)
+        .contiguous()
+    )
+    if pad_N > 0 or pad_K > 0:
+        weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
+        weight_scaled = weight_scaled[..., :N, :K].contiguous()
+    else:
+        weight_scaled = weight_scaled.view(E, N, K)
+    return weight_scaled
+
+
+def ref_w8a16_block_fp8_moe(
+    a: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_s: torch.Tensor,
+    w2_s: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    block_size: list[int],
+) -> torch.Tensor:
+    """Reference FP8 W8A16 block-scaled fused MoE in pure torch."""
+    B, D = a.shape
+    topk = topk_ids.size(1)
+
+    w1_dq = _block_dequant_weight(w1, w1_s, block_size)
+    w2_dq = _block_dequant_weight(w2, w2_s, block_size)
+
+    a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
+    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
+
+    topk_weight_flat = topk_weight.view(-1)
+    topk_ids_flat = topk_ids.view(-1)
+
+    for i in range(w1_dq.shape[0]):
+        mask = topk_ids_flat == i
+        if mask.sum():
+            ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
+            ic1 = _silu_and_mul(ic0)
+            out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
+
+    return (
+        (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
+        .sum(dim=1)
+        .to(a.dtype)
+    )
+
+
+def _make_fp8_moe_weights(
+    E: int,
+    N: int,
+    K: int,
+    block_size: list[int],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Generate random FP8 MoE weights with random block scales."""
+    block_n, block_k = block_size
+
+    w1 = (
+        (torch.randn(E, 2 * N, K) * FP8_SCALE)
+        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
+        .to(torch.float8_e4m3fn)
+    )
+    w2 = (
+        (torch.randn(E, K, N) * FP8_SCALE)
+        .clamp(min=-FP8_SCALE, max=FP8_SCALE)
+        .to(torch.float8_e4m3fn)
+    )
+
+    w1_s = (
+        torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
+        * FACTOR_FOR_SCALE
+    )
+    w2_s = (
+        torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
+        * FACTOR_FOR_SCALE
+    )
+    return w1, w2, w1_s, w2_s
+
+
+FP8_NUM_TOKENS = [1, 2, 64, 121]
+FP8_MOE_CONFIGS = [
+    (256, 512, 8, 2),
+    (256, 512, 8, 4),
+    (512, 256, 8, 2),
+    (512, 256, 8, 4),
+    (512, 512, 8, 2),
+    (512, 512, 8, 4),
+    (768, 2048, 8, 2),
+    (768, 2048, 8, 4),
+    (768, 2048, 128, 8),
+]
+
+
+@pytest.mark.parametrize("M", FP8_NUM_TOKENS)
+@pytest.mark.parametrize("N,K,E,topk", FP8_MOE_CONFIGS)
+@pytest.mark.parametrize("seed", [0])
+def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
+    """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
+    set_random_seed(seed)
+
+    a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
+    w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
+
+    score = torch.randn(M, E, dtype=torch.bfloat16)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_ids = topk_ids.to(torch.int32)
+
+    ref_out = ref_w8a16_block_fp8_moe(
+        a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
+    )
+
+    pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
+
+    # Test inplace=False against reference
+    out = ops.fused_experts_cpu(
+        a.clone(),
+        pw1,
+        pw2,
+        topk_weight,
+        topk_ids,
+        False,
+        ops.CPUQuantMethod.FP8_W8A16,
+        w1_s,
+        w2_s,
+        None,
+        None,
+        BLOCK_SIZE,
+        is_vnni=True,
+    )
+    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
+
+    # Test inplace=True produces identical output
+    out_inplace = ops.fused_experts_cpu(
+        a.clone(),
+        pw1,
+        pw2,
+        topk_weight,
+        topk_ids,
+        True,
+        ops.CPUQuantMethod.FP8_W8A16,
+        w1_s,
+        w2_s,
+        None,
+        None,
+        BLOCK_SIZE,
+        is_vnni=True,
+    )
+    torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
+
+
+# MXFP4 W4A16 fused MoE
+
+
+class MXFP4QuantizeUtil:
+    """MXFP4 quantization utility."""
+
+    E2M1_max = 6.0
+    E2M1_values = [0, 0.5, 1, 1.5, 2, 3, 4, 6]
+    E2M1_bounds = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5])
+    block_size = 32
+
+    @classmethod
+    def quantize(cls, input: torch.Tensor) -> tuple:
+        """Quantize BF16 tensor to MXFP4 packed uint8 format."""
+
+        def cast_fp4(x):
+            sign = torch.sign(x)
+            sign_bit = (2 - sign) // 2
+            ord_ = torch.sum(
+                (x.abs().unsqueeze(-1) - cls.E2M1_bounds.to(x.device)) > 0, dim=-1
+            )
+            fp4_val = (sign_bit * 0b1000 + ord_).to(torch.uint8)
+            return fp4_val
+
+        def fuse_uint4_to_uint8(x):
+            left_side = x[..., 0::2]
+            right_side = x[..., 1::2]
+            new_data = right_side.clone() << 4
+            new_data[..., : left_side.shape[-1]] += left_side
+            return new_data
+
+        original_shape = input.shape
+        input = input.view(-1, cls.block_size)
+        input_amax = input.abs().max(dim=-1, keepdim=True).values
+        descale = input_amax / cls.E2M1_max
+        min_value = torch.tensor(-127.0, device=descale.device)
+        e8m0_scale = torch.ceil(torch.maximum(torch.log2(descale), min_value))
+
+        input = (input / torch.exp2(e8m0_scale)).view(original_shape)
+        input_q = cast_fp4(input)
+        input_q = fuse_uint4_to_uint8(input_q)
+        e8m0_scale = (e8m0_scale + 127).to(torch.uint8)
+        return input_q, e8m0_scale
+
+    @classmethod
+    def dequantize(cls, quantized_data, dtype: torch.dtype, scale):
+        """Dequantize MXFP4 packed tensor back to float."""
+
+        def unfuse_uint8_to_uint4(x):
+            left_side = x & 0x0F
+            right_side = (x >> 4) & 0x0F
+            shape = list(x.shape)
+            shape[-1] = shape[-1] * 2
+            result = torch.zeros(shape, dtype=torch.uint8, device=x.device)
+            result[..., 0::2] = left_side
+            result[..., 1::2] = right_side
+            return result
+
+        e8m0_scale = scale
+        x_unfused = unfuse_uint8_to_uint4(quantized_data)
+        sign = 1 - 2 * ((x_unfused & 0b1000) >> 3).to(torch.float32)
+        magnitude = (x_unfused & 0b0111).to(torch.long)
+        values = torch.tensor(cls.E2M1_values, device=quantized_data.device)
+        original_shape = magnitude.shape
+        x_float = values[magnitude.reshape(-1)].reshape(original_shape)
+        x_float = sign.float() * x_float
+        x_float = x_float.reshape(-1, cls.block_size)
+        scale_factor = torch.exp2(e8m0_scale.float() - 127)
+        scale_factor = scale_factor.reshape(-1, 1)
+        x_float = x_float * scale_factor
+        return x_float.reshape(original_shape).to(dtype)
+
+
+def _swiglu(x: torch.Tensor, alpha: float, limit: float) -> torch.Tensor:
+    """SwigLU activation used in GPT-OSS.
+    Input is interleaved: [gate_0, up_0, gate_1, up_1, ...] in last dim.
+    """
+    gate = x[..., 0::2]
+    up = x[..., 1::2]
+    gate_clamped = torch.clamp(gate, max=limit)
+    up_clamped = torch.clamp(up, min=-limit, max=limit)
+    return gate_clamped * torch.sigmoid(alpha * gate_clamped) * (up_clamped + 1)
+
+
+def ref_mxfp4_fused_moe(
+    a: torch.Tensor,
+    w1_dq: torch.Tensor,
+    w2_dq: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    """Reference MXFP4 fused MoE with SiLU activation."""
+    B, D = a.shape
+    a_f = a.float()
+    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
+    topk_ids_flat = topk_ids.view(-1)
+
+    for i in range(w1_dq.shape[0]):
+        mask = topk_ids_flat == i
+        if mask.sum() == 0:
+            continue
+        token_indices = torch.where(mask)[0]
+        source_indices = token_indices // topk
+        ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
+        ic1 = _silu_and_mul(ic0)
+        out[mask] = torch.matmul(ic1, w2_dq[i].float().T)
+
+    return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
+
+
+def ref_mxfp4_fused_moe_gptoss(
+    a: torch.Tensor,
+    w1_dq: torch.Tensor,
+    w2_dq: torch.Tensor,
+    w1_bias: torch.Tensor,
+    w2_bias: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+    alpha: float,
+    limit: float,
+) -> torch.Tensor:
+    """Reference MXFP4 fused MoE with SwigLU+bias (GPT-OSS style)."""
+    B, D = a.shape
+    topk = topk_ids.shape[1]
+    a_f = a.float()
+    E = w1_dq.shape[0]
+    out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
+    topk_ids_flat = topk_ids.view(-1)
+
+    for i in range(E):
+        mask = topk_ids_flat == i
+        if mask.sum() == 0:
+            continue
+        token_indices = torch.where(mask)[0]
+        source_indices = token_indices // topk
+        ic0 = torch.matmul(a_f[source_indices], w1_dq[i].float().T)
+        ic0 = ic0 + w1_bias[i].float()
+        ic1 = _swiglu(ic0, alpha, limit)
+        ic2 = torch.matmul(ic1, w2_dq[i].float().T)
+        ic2 = ic2 + w2_bias[i].float()
+        out[mask] = ic2
+
+    return (out.view(B, topk, -1) * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
+
+
+def _prepack_mxfp4_experts(
+    w: torch.Tensor, w_scale: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """VNNI-prepack MXFP4 weights and repack scales."""
+    packed_w = torch.ops._C.convert_weight_packed(w)
+    packed_s = torch.ops._C.convert_scale_packed(w_scale)
+    return packed_w, packed_s
+
+
+MXFP4_NUM_TOKENS = [1, 2, 32, 121]
+MXFP4_MOE_CONFIGS = [
+    (128, 128, 4, 2),
+    (256, 256, 8, 4),
+    (352, 256, 8, 4),
+    (512, 320, 8, 4),
+]
+
+
+@pytest.mark.parametrize("M", MXFP4_NUM_TOKENS)
+@pytest.mark.parametrize("N,K,E,topk", MXFP4_MOE_CONFIGS)
+@pytest.mark.parametrize("seed", [0])
+def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
+    """Test fused_experts_mxfp4_cpu against dequantized torch reference."""
+    set_random_seed(seed)
+    dtype = torch.bfloat16
+
+    a = torch.randn(M, K, dtype=dtype) / 10
+
+    # Generate and quantize weights
+    w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
+    w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
+    w1s = w1s.reshape(E, 2 * N, K // 32)
+    w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
+
+    w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
+    w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
+    w2s = w2s.reshape(E, K, N // 32)
+    w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
+
+    # Routing
+    score = torch.randn(M, E, dtype=dtype)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_ids = topk_ids.to(torch.int32)
+
+    # Reference
+    ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, topk)
+
+    # Pack weights for kernel
+    pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
+    pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
+
+    # Kernel
+    out = ops.fused_experts_cpu(
+        a.clone(),
+        pw1,
+        pw2,
+        topk_weight,
+        topk_ids,
+        False,  # inplace
+        ops.CPUQuantMethod.MXFP4,
+        pw1s,  # w1_scale
+        pw2s,  # w2_scale
+        None,  # w1_zero
+        None,  # w2_zero
+        None,  # block_size
+    )
+
+    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.parametrize("M", [1, 32])
+@pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)])
+@pytest.mark.parametrize("seed", [0])
+def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
+    """Test fused_experts_mxfp4_cpu with bias and SwigLU activation (GPT-OSS)."""
+    set_random_seed(seed)
+    dtype = torch.bfloat16
+    alpha = 1.702
+    limit = 7.0
+
+    a = torch.randn(M, K, dtype=dtype) / 10
+
+    # Generate and quantize weights
+    w1_bf16 = torch.randn(E, 2 * N, K, dtype=dtype) / 10
+    w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
+    w1s = w1s.reshape(E, 2 * N, K // 32)
+    w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
+    w1_b = torch.randn(E, 2 * N, dtype=torch.float32) / 10
+
+    w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10
+    w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
+    w2s = w2s.reshape(E, K, N // 32)
+    w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
+    w2_b = torch.randn(E, K, dtype=torch.float32) / 10
+
+    # Routing
+    score = torch.randn(M, E, dtype=dtype)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_ids = topk_ids.to(torch.int32)
+
+    # Reference
+    ref_out = ref_mxfp4_fused_moe_gptoss(
+        a, w1dq, w2dq, w1_b, w2_b, topk_weight, topk_ids, alpha, limit
+    )
+
+    # Pack weights for kernel
+    pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
+    pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
+
+    # Kernel
+    out = ops.fused_experts_cpu(
+        a.clone(),
+        pw1,
+        pw2,
+        topk_weight,
+        topk_ids,
+        False,  # inplace
+        ops.CPUQuantMethod.MXFP4,
+        pw1s,  # w1_scale
+        pw2s,  # w2_scale
+        None,  # w1_zero
+        None,  # w2_zero
+        None,  # block_size
+        w1_bias=w1_b,
+        w2_bias=w2_b,
+        alpha=alpha,
+        limit=limit,
+    )
+
+    torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))
@@ -453,11 +453,15 @@ def test_flashinfer_cutedsl_moe_masked(
    )

    # reference
-    a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, input_global_scale)
+    # input_global_scale is per-expert ([num_experts]); fp4_quantize and
+    # dequantize_nvfp4_to_dtype are non-grouped APIs that expect [1] or
+    # [num_tokens]. Use a single element since all values are uniform here.
+    a_global = input_global_scale[:1].contiguous()
+    a_fp4, a_scale_interleaved = fp4_quantize(hidden_states, a_global)
    a_in_dtype = dequantize_nvfp4_to_dtype(
        a_fp4,
        a_scale_interleaved,
-        input_global_scale,
+        a_global,
        dtype=hidden_states.dtype,
        device=hidden_states.device,
        block_size=16,
@@ -746,7 +746,8 @@ def marlin_moe_generate_valid_test_cases():
        for sub_case in inner_combinations:
            if (
                sub_case[0] == scalar_types.float8_e4m3fn
-                and current_platform.get_device_capability() not in [89, 120]
+                and not current_platform.is_device_capability(89)
+                and not current_platform.is_device_capability_family(120)
            ):
                continue

@@ -897,6 +898,7 @@ class MarlinMoEWeightData:
    marlin_moe_generate_valid_test_cases(),
 )
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
+@pytest.mark.usefixtures("default_vllm_config")
 def test_fused_marlin_moe(
    a_type: ScalarType,
    b_type: ScalarType,
@@ -1009,6 +1011,7 @@ def test_fused_marlin_moe(

@pytest.mark.flaky(reruns=2)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
+@pytest.mark.usefixtures("default_vllm_config")
@pytest.mark.parametrize("m", [1, 256])
 def test_fused_marlin_moe_with_bias(m):
    set_random_seed(0)
@@ -1081,6 +1084,7 @@ def test_fused_marlin_moe_with_bias(m):

@pytest.mark.flaky(reruns=2)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
+@pytest.mark.usefixtures("default_vllm_config")
@pytest.mark.parametrize("m", [1, 64, 256])
@pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])
@pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])
@@ -1334,36 +1338,96 @@ def test_cpu_fused_moe_basic(
    torch.testing.assert_close(out, ref, atol=atol, rtol=0)


-@pytest.mark.parametrize("m", [16, 32, 64])
-@pytest.mark.parametrize("n", [128])
-@pytest.mark.parametrize("k", [128])
-@pytest.mark.parametrize("e", [8, 12, 16, 32])
-@pytest.mark.parametrize("topk", [2, 4])
-@pytest.mark.parametrize("max_tokens_per_batch", [16, 32, 64])
+def _batched_fused_marlin_moe_cases() -> list[Any]:
+    cases = [
+        pytest.param(
+            m,
+            128,
+            128,
+            e,
+            topk,
+            max_tokens_per_batch,
+            torch.bfloat16,
+            scalar_types.float4_e2m1f,
+            None,
+            1e-3,
+            id=(
+                f"m{m}-n128-k128-e{e}-topk{topk}-max_tokens{max_tokens_per_batch}-mxfp4"
+            ),
+        )
+        for m in [16, 32, 64]
+        for e in [8, 12, 16, 32]
+        for topk in [2, 4]
+        for max_tokens_per_batch in [16, 32, 64]
+    ]
+    cases.append(
+        pytest.param(
+            32,
+            128,
+            128,
+            8,
+            2,
+            64,
+            torch.float16,
+            scalar_types.uint4,
+            scalar_types.int8,
+            4e-2,
+            id="awq-int8-activation-metadata",
+        )
+    )
+    return cases
+
+
+@pytest.mark.parametrize(
+    ("m,n,k,e,topk,max_tokens_per_batch,dtype,quant_dtype,input_type,atol"),
+    _batched_fused_marlin_moe_cases(),
+)
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
 def test_batched_fused_marlin_moe(
-    m: int, n: int, k: int, e: int, topk: int, max_tokens_per_batch: int
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    max_tokens_per_batch: int,
+    dtype: torch.dtype,
+    quant_dtype: ScalarType,
+    input_type: ScalarType | None,
+    atol: float,
 ):
    print(
        f"testing m={m}, n={n}, k={k}, e={e}, "
        f"topk={topk}, "
-        f"max_tokens_per_batch={max_tokens_per_batch}"
+        f"max_tokens_per_batch={max_tokens_per_batch}, "
+        f"dtype={dtype}, quant_dtype={quant_dtype}, input_type={input_type}"
    )
    set_random_seed(0)

-    dtype = torch.bfloat16
-    quant_dtype = scalar_types.float4_e2m1f
    group_size = 32
+    if input_type == scalar_types.int8:
+        input_dtype = torch.int8
+    elif input_type == scalar_types.float8_e4m3fn:
+        input_dtype = torch.float8_e4m3fn
+    else:
+        input_dtype = None

    a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
    w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
    w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20

    w1_data = MarlinMoEWeightData.make(
-        w=w1, quant_type=quant_dtype, group_size=group_size, act_order=None
+        w=w1,
+        quant_type=quant_dtype,
+        group_size=group_size,
+        act_order=None,
+        input_type=input_type,
    )
    w2_data = MarlinMoEWeightData.make(
-        w=w2, quant_type=quant_dtype, group_size=group_size, act_order=None
+        w=w2,
+        quant_type=quant_dtype,
+        group_size=group_size,
+        act_order=None,
+        input_type=input_type,
    )

    score = torch.randn((m, e), device="cuda", dtype=dtype)
@@ -1483,6 +1547,12 @@ def test_batched_fused_marlin_moe(
        "quant_type_id": quant_dtype.id,
        "is_k_full": True,
    }
+    if input_dtype is not None:
+        kwargs["input_dtype"] = input_dtype
+        if w1_data.a_scales_factor is not None:
+            kwargs["input_global_scale1"] = w1_data.a_scales_factor
+        if w2_data.a_scales_factor is not None:
+            kwargs["input_global_scale2"] = w2_data.a_scales_factor

    # Reference
    fused_marlin_moe_kwargs = kwargs | {
@@ -1498,7 +1568,7 @@ def test_batched_fused_marlin_moe(
        pytest.skip("Cannot represent data in Batched Format.")
    marlin_output = br.run(a, kwargs)

-    torch.testing.assert_close(marlin_output, ref_marlin_output, atol=1e-3, rtol=0)
+    torch.testing.assert_close(marlin_output, ref_marlin_output, atol=atol, rtol=0)


@pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)])
@@ -381,7 +381,8 @@ def marlin_generate_valid_test_cases():
        for sub_case in inner_combinations:
            if (
                sub_case[0] == scalar_types.float8_e4m3fn
-                and current_platform.get_device_capability() not in [89, 120]
+                and not current_platform.is_device_capability(89)
+                and not current_platform.is_device_capability_family(120)
            ):
                continue
            args = sub_case + (size_m, size_n, size_k) + case[4:]
@@ -313,6 +313,16 @@ def qwen35_vl_lora_files():
    return snapshot_download(repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora")


+@pytest.fixture(scope="session")
+def qwen36_moe_2d_lora_files():
+    return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora")
+
+
+@pytest.fixture(scope="session")
+def qwen36_moe_3d_lora_files():
+    return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora")
+
+
@pytest.fixture
 def reset_default_device():
    """
@@ -179,6 +179,100 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):
    )


+def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_model):
+    """A module reachable from two attribute paths (e.g. a MoE gate held
+    both directly on the block and inside its inner runner) must produce a
+    single LoRA wrapper. Both paths must end up pointing to that same
+    wrapper instance, and only the canonical path should live in
+    `manager.modules` — otherwise activate_adapter would call `reset_lora`
+    on the alias and clobber weights set under the canonical name.
+    """
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    class AliasContainer(nn.Module):
+        def __init__(self, gate: nn.Module):
+            super().__init__()
+            self.gate = gate  # canonical path: "moe.gate"
+
+            # Inner submodule holding the SAME gate instance under another
+            # path. This mirrors how FusedMoE.runner.gate references the
+            # block's gate in qwen3_moe.
+            class _Runner(nn.Module):
+                def __init__(self, g):
+                    super().__init__()
+                    self.gate = g  # alias path: "moe.runner.gate"
+
+            self.runner = _Runner(gate)
+
+    gate = ReplicatedLinear(10, 4, bias=False)
+    model = dummy_model
+    model.add_module("moe", AliasContainer(gate))
+
+    assert model.moe.gate is model.moe.runner.gate
+
+    manager = LoRAModelManager(
+        model,
+        1,
+        1,
+        1,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
+        ),
+        torch.device(DEVICES[0]),
+    )
+
+    canonical = manager.model.get_submodule("moe.gate")
+    alias = manager.model.get_submodule("moe.runner.gate")
+
+    # Same wrapper instance on both paths so forward through either side
+    # sees the LoRA-augmented module.
+    assert isinstance(canonical, ReplicatedLinearWithLoRA)
+    assert alias is canonical
+
+    # Only the canonical path is tracked as a LoRA target. Tracking the
+    # alias would cause activate_adapter to reset_lora on it after the
+    # canonical entry already populated the weights.
+    assert "moe.gate" in manager.modules
+    assert "moe.runner.gate" not in manager.modules
+
+
+def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):
+    """The dedup logic must NOT collapse `lm_head` even when it is reachable
+    from another attribute path (tied-embedding models do
+    `self.lm_head = self.model.embed_tokens`, sharing the same nn.Module
+    instance). The lm_head branch additionally rewires `logits_processor`
+    into a `LogitsProcessorWithLoRA`, so skipping it would silently break
+    LoRA on lm_head.
+    """
+    from vllm.lora.layers import LogitsProcessorWithLoRA
+
+    # Add a non-lm_head alias to the same module instance as lm_head. The
+    # dedup keys on id(module); without the lm_head exemption the alias
+    # would consume the wrapped_by_id slot first and lm_head would be
+    # silently skipped, so logits_processor would never be wrapped.
+    model = dummy_model
+    model.add_module("embed_tokens", model.lm_head)
+    assert model.embed_tokens is model.lm_head
+
+    manager = LoRAModelManager(
+        model,
+        1,
+        1,
+        1,
+        LoRAConfig(
+            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
+        ),
+        torch.device(DEVICES[0]),
+    )
+
+    # lm_head's special handling still ran: logits_processor got wrapped
+    # and the lm_head entry is tracked under self.modules.
+    assert isinstance(
+        manager.model.get_submodule("logits_processor"), LogitsProcessorWithLoRA
+    )
+    assert "lm_head" in manager.modules
+
+
 def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model):
    class UnsupportedContainer(nn.Module):
        def __init__(self):
@@ -104,5 +104,222 @@ def test_moe_lora_align_block_size(
                assert torch.all(topk_ids.view(-1)[indices] == expert_id)


+# Sentinel values for the regression tests below. Distinctive out-of-domain
+# ints so that "kernel never wrote this slot" is directly observable: the
+# kernel only ever writes a real expert id in [0, num_experts) or -1
+# (expert_ids), a token index or the `numel` padding value (sorted_token_ids),
+# and a block-aligned cumsum count (num_tokens_post_pad).
+SENTINEL_EXPERT = -2
+SENTINEL_TOKEN = -7
+SENTINEL_NPAD = -13
+
+
+def _build_and_run_align(
+    *,
+    num_lora_tokens,
+    num_base_tokens,
+    max_loras,
+    num_experts=64,
+    topk_num=6,
+    block_size=16,
+    lora_ids_override=None,
+    disabled_slots=(),
+    seed=1,
+):
+    """Build inputs the way ``LoRAKernelMeta.prepare_tensors`` does, run
+    ``moe_lora_align_block_size``, and return a dict of result tensors plus
+    derived sizes. Output buffers are pre-filled with ``SENTINEL_*`` so
+    callers can assert which slots the kernel did / did not touch.
+
+    Tokens are assigned to LoRA slot 0 (first ``num_lora_tokens``) then -1
+    (remaining ``num_base_tokens``), matching the "mixed base + 1 LoRA"
+    shape used to repro vllm-project/vllm#32235.
+
+    ``lora_ids_override``: optional 1-D int tensor of length ``max_loras+1``
+    used verbatim. Default mirrors ``prepare_tensors`` (sorted-unique into
+    the head, -1 tail).
+    ``disabled_slots``: iterable of slot indices to clear in ``adapter_enabled``.
+    """
+    random.seed(seed)
+    num_tokens = num_lora_tokens + num_base_tokens
+    assert num_tokens > 0, "test requires at least one token"
+
+    topk_ids = torch.zeros((num_tokens, topk_num), dtype=torch.int32)
+    token_lora_mapping = torch.empty((num_tokens,), dtype=torch.int32)
+    for i in range(num_tokens):
+        pool = list(range(num_experts))
+        random.shuffle(pool)
+        for j in range(topk_num):
+            topk_ids[i, j] = pool[j]
+        token_lora_mapping[i] = 0 if i < num_lora_tokens else -1
+    topk_ids = topk_ids.to(DEVICE_TYPE)
+    token_lora_mapping = token_lora_mapping.to(DEVICE_TYPE)
+
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    if topk_ids.numel() < num_experts:
+        max_num_tokens_padded = topk_ids.numel() * block_size
+    max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
+
+    if lora_ids_override is None:
+        lora_ids = torch.full(
+            (max_loras + 1,), -1, dtype=torch.int32, device=DEVICE_TYPE
+        )
+        unique_ids = torch.unique(token_lora_mapping, sorted=True)
+        lora_ids[: unique_ids.numel()] = unique_ids.to(torch.int32)
+    else:
+        assert lora_ids_override.numel() == max_loras + 1
+        lora_ids = lora_ids_override.to(dtype=torch.int32, device=DEVICE_TYPE)
+
+    adapter_enabled = torch.ones(
+        (max_loras + 1,), dtype=torch.int32, device=DEVICE_TYPE
+    )
+    for slot in disabled_slots:
+        adapter_enabled[slot] = 0
+
+    sorted_token_ids = torch.full(
+        (max_loras * max_num_tokens_padded,),
+        SENTINEL_TOKEN,
+        dtype=torch.int32,
+        device=DEVICE_TYPE,
+    )
+    expert_ids = torch.full(
+        (max_loras * max_num_m_blocks,),
+        SENTINEL_EXPERT,
+        dtype=torch.int32,
+        device=DEVICE_TYPE,
+    )
+    num_tokens_post_pad = torch.full(
+        (max_loras,), SENTINEL_NPAD, dtype=torch.int32, device=DEVICE_TYPE
+    )
+
+    ops.moe_lora_align_block_size(
+        topk_ids,
+        token_lora_mapping,
+        num_experts,
+        block_size,
+        max_loras,
+        max_num_tokens_padded,
+        max_num_m_blocks,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        adapter_enabled,
+        lora_ids,
+    )
+
+    return {
+        "lora_ids": lora_ids,
+        "sorted_token_ids": sorted_token_ids,
+        "expert_ids": expert_ids,
+        "num_tokens_post_pad": num_tokens_post_pad,
+        "max_num_tokens_padded": max_num_tokens_padded,
+        "block_size": block_size,
+        "max_loras": max_loras,
+    }
+
+
+@pytest.mark.parametrize(
+    "max_loras",
+    [
+        1,
+        2,
+    ],
+)
+def test_moe_lora_align_block_size_mixed_base_and_lora(max_loras):
+    """Regression test for issue #32235: real LoRA slot must not be skipped
+    when ``active_lora_ids`` has -1 at position 0."""
+    out = _build_and_run_align(
+        num_lora_tokens=8, num_base_tokens=8, max_loras=max_loras
+    )
+
+    # Sanity check on the layout being tested.
+    assert out["lora_ids"][0].item() == -1, (
+        "prepare_tensors layout mismatch: -1 expected at position 0 for mixed batch"
+    )
+
+    real_slot = 0
+    post_pad = out["num_tokens_post_pad"][real_slot].item()
+    assert post_pad != SENTINEL_NPAD, (
+        f"num_tokens_post_pad[{real_slot}] was never written by the kernel; "
+        "the align kernel skipped the real LoRA slot."
+    )
+    assert (
+        0 < post_pad <= out["max_num_tokens_padded"]
+        and post_pad % out["block_size"] == 0
+    ), f"num_tokens_post_pad[{real_slot}]={post_pad} is not a valid block-aligned count"
+
+    expert_row = out["expert_ids"].view(max_loras, -1)[real_slot]
+    assert (expert_row != SENTINEL_EXPERT).all(), (
+        f"expert_ids row for slot {real_slot} has unwritten sentinel entries; "
+        "the align kernel skipped the real LoRA slot."
+    )
+
+    sorted_row = out["sorted_token_ids"].view(max_loras, -1)[real_slot]
+    assert (sorted_row != SENTINEL_TOKEN).all(), (
+        f"sorted_token_ids row for slot {real_slot} has unwritten sentinel "
+        "entries; the align kernel skipped the real LoRA slot."
+    )
+
+
+def test_moe_lora_align_block_size_disabled_adapter_untouched():
+    """Disabled-adapter slot rows must remain untouched by all three align
+    kernels. Pins the invariant protected by the ``adapter_enabled`` guard
+    in ``lora_count_and_sort_expert_tokens_kernel``: without it the sort
+    kernel reads uninitialized ``token_mask`` values for disabled slots and
+    pollutes ``sorted_token_ids`` / ``cumsum_buffer``."""
+    max_loras = 1
+    out = _build_and_run_align(
+        num_lora_tokens=16,
+        num_base_tokens=0,
+        max_loras=max_loras,
+        disabled_slots=(0,),
+    )
+    # Sanity: slot 0 IS present in active_lora_ids (otherwise we would only
+    # exercise the lora_id == -1 / >= max_loras guards).
+    assert (out["lora_ids"] == 0).any().item()
+
+    assert out["num_tokens_post_pad"][0].item() == SENTINEL_NPAD, (
+        "num_tokens_post_pad[0] was modified for a disabled adapter slot."
+    )
+    expert_row = out["expert_ids"].view(max_loras, -1)[0]
+    assert (expert_row == SENTINEL_EXPERT).all(), (
+        "expert_ids row for disabled slot 0 was partially written."
+    )
+    # Row specifically protected by the sort-kernel adapter_enabled guard.
+    sorted_row = out["sorted_token_ids"].view(max_loras, -1)[0]
+    assert (sorted_row == SENTINEL_TOKEN).all(), (
+        "sorted_token_ids row for disabled slot 0 was polluted by the sort "
+        "kernel; lora_count_and_sort_expert_tokens_kernel must skip "
+        "adapter_enabled == 0 slots."
+    )
+
+
+def test_moe_lora_align_block_size_lora_id_oob_guard():
+    """Regression test for the ``lora_id >= max_loras`` guard.
+
+    Production ``LoRAKernelMeta.prepare_tensors`` pre-fills the tail of
+    ``active_lora_ids`` with -1, so the existing ``lora_id == -1`` check
+    covers the extra slot. This test bypasses that invariant and injects
+    an out-of-range value (5 with max_loras=1) at the tail to verify the
+    explicit guard prevents OOB reads against ``adapter_enabled`` and
+    OOB writes against the max_loras-sized output buffers. Without the
+    guard, an illegal-memory-access would surface on the next CUDA sync.
+    """
+    max_loras = 1
+    lora_ids_override = torch.tensor([0, 5], dtype=torch.int32)
+    out = _build_and_run_align(
+        num_lora_tokens=16,
+        num_base_tokens=0,
+        max_loras=max_loras,
+        lora_ids_override=lora_ids_override,
+    )
+    # The .item() call below syncs and would surface any async
+    # illegal-memory-access from the OOB iteration.
+    assert out["num_tokens_post_pad"][0].item() != SENTINEL_NPAD, (
+        "real LoRA slot 0 was skipped by the align kernel"
+    )
+
+
 if __name__ == "__main__":
    pytest.main([__file__])
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+import vllm
+import vllm.config
+from vllm.assets.image import ImageAsset
+from vllm.lora.request import LoRARequest
+
+from ..utils import multi_gpu_test
+
+MODEL_PATH = "Qwen/Qwen3.6-35B-A3B"
+
+LORA_2D_ID = 1
+LORA_3D_ID = 2
+
+PROMPT_TEMPLATE = """<|im_start|>user
+<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>
+<|im_start|>assistant
+<think>
+
+</think>
+"""
+
+# Visual captioning prompts: each image will be paired with one LoRA in the
+# mixed-batch case so we can check per-prompt routing.
+VL_TEST_IMAGES = [
+    ImageAsset("stop_sign"),
+    ImageAsset("cherry_blossom"),
+]
+
+
+def _build_prompts() -> list[dict]:
+    return [
+        {
+            "prompt": PROMPT_TEMPLATE,
+            "multi_modal_data": {"image": asset.pil_image},
+        }
+        for asset in VL_TEST_IMAGES
+    ]
+
+
+def _generate(llm: vllm.LLM, lora_request) -> list[str]:
+    outputs = llm.generate(
+        _build_prompts(),
+        vllm.SamplingParams(temperature=0, max_tokens=128),
+        lora_request=lora_request,
+    )
+    return [out.outputs[0].text.strip() for out in outputs]
+
+
+def _run_mixed_2d_3d_lora_test(
+    lora_2d_files: str,
+    lora_3d_files: str,
+    tensor_parallel_size: int,
+    fully_sharded_loras: bool,
+) -> None:
+    llm = vllm.LLM(
+        model=MODEL_PATH,
+        max_model_len=4096,
+        enable_lora=True,
+        enable_mixed_moe_lora_format=True,
+        max_loras=2,
+        max_lora_rank=8,
+        max_num_seqs=4,
+        enforce_eager=True,
+        tensor_parallel_size=tensor_parallel_size,
+        enable_expert_parallel=not fully_sharded_loras,
+        fully_sharded_loras=fully_sharded_loras,
+        trust_remote_code=True,
+        enable_tower_connector_lora=True,
+        mm_processor_cache_gb=0,
+        limit_mm_per_prompt={"image": 1},
+        compilation_config=vllm.config.CompilationConfig(
+            cudagraph_specialize_lora=False,
+        ),
+    )
+
+    lora_2d = LoRARequest(
+        "lora_2d",
+        LORA_2D_ID,
+        lora_2d_files,
+        is_3d_lora_weight=False,
+    )
+    lora_3d = LoRARequest(
+        "lora_3d",
+        LORA_3D_ID,
+        lora_3d_files,
+        is_3d_lora_weight=True,
+    )
+
+    # Reference: each adapter alone over both prompts.
+    outputs_2d_alone = _generate(llm, lora_2d)
+    outputs_3d_alone = _generate(llm, lora_3d)
+
+    assert len(outputs_2d_alone) == len(VL_TEST_IMAGES)
+    assert len(outputs_3d_alone) == len(VL_TEST_IMAGES)
+    for text in outputs_2d_alone + outputs_3d_alone:
+        assert text, "Empty output from single-adapter LoRA generation"
+
+    # Mixed batch: prompt 0 uses the 2D adapter, prompt 1 uses the 3D
+    # adapter. Per-prompt outputs must match the standalone runs.
+    mixed_outputs = _generate(llm, [lora_2d, lora_3d])
+
+    assert mixed_outputs[0] == outputs_2d_alone[0], (
+        f"Mixed-batch 2D output {mixed_outputs[0]!r} does not match "
+        f"standalone 2D output {outputs_2d_alone[0]!r}"
+    )
+    assert mixed_outputs[1] == outputs_3d_alone[1], (
+        f"Mixed-batch 3D output {mixed_outputs[1]!r} does not match "
+        f"standalone 3D output {outputs_3d_alone[1]!r}"
+    )
+
+    # Reverse assignment: neither adapter should be silently aliased.
+    swapped_outputs = _generate(llm, [lora_3d, lora_2d])
+    assert swapped_outputs[0] == outputs_3d_alone[0], (
+        f"Swapped-batch 3D output {swapped_outputs[0]!r} does not match "
+        f"standalone 3D output {outputs_3d_alone[0]!r}"
+    )
+    assert swapped_outputs[1] == outputs_2d_alone[1], (
+        f"Swapped-batch 2D output {swapped_outputs[1]!r} does not match "
+        f"standalone 2D output {outputs_2d_alone[1]!r}"
+    )
+
+
+@pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+@multi_gpu_test(num_gpus=2)
+def test_qwen36_moe_mixed_2d_3d_lora_tp2(
+    qwen36_moe_2d_lora_files,
+    qwen36_moe_3d_lora_files,
+    fully_sharded_loras,
+):
+    _run_mixed_2d_3d_lora_test(
+        lora_2d_files=qwen36_moe_2d_lora_files,
+        lora_3d_files=qwen36_moe_3d_lora_files,
+        tensor_parallel_size=2,
+        fully_sharded_loras=fully_sharded_loras,
+    )
+
+
+@pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
+@multi_gpu_test(num_gpus=4)
+def test_qwen36_moe_mixed_2d_3d_lora_tp4(
+    qwen36_moe_2d_lora_files,
+    qwen36_moe_3d_lora_files,
+    fully_sharded_loras,
+):
+    _run_mixed_2d_3d_lora_test(
+        lora_2d_files=qwen36_moe_2d_lora_files,
+        lora_3d_files=qwen36_moe_3d_lora_files,
+        tensor_parallel_size=4,
+        fully_sharded_loras=fully_sharded_loras,
+    )
@@ -77,6 +77,7 @@ def test_qwen3moe_lora(qwen3moe_lora_files):
        enforce_eager=True,
        trust_remote_code=True,
        enable_chunked_prefill=True,
+        enable_mixed_moe_lora_format=True,
    )

    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
@@ -85,7 +86,8 @@ def test_qwen3moe_lora(qwen3moe_lora_files):

@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize("ep", [False, True])
-def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):
+@pytest.mark.parametrize("enable_mixed_moe_lora_format", [False, True])
+def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files, enable_mixed_moe_lora_format):
    llm = vllm.LLM(
        MODEL_PATH,
        max_model_len=1024,
@@ -94,6 +96,7 @@ def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):
        trust_remote_code=True,
        enable_expert_parallel=ep,
        tensor_parallel_size=2,
+        enable_mixed_moe_lora_format=enable_mixed_moe_lora_format,
    )

    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
@@ -102,7 +105,8 @@ def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):

@multi_gpu_test(num_gpus=4)
@pytest.mark.parametrize("ep", [False, True])
-def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files):
+@pytest.mark.parametrize("enable_mixed_moe_lora_format", [False, True])
+def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files, enable_mixed_moe_lora_format):
    llm = vllm.LLM(
        MODEL_PATH,
        max_model_len=1024,
@@ -111,6 +115,7 @@ def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files):
        trust_remote_code=True,
        enable_expert_parallel=ep,
        tensor_parallel_size=4,
+        enable_mixed_moe_lora_format=enable_mixed_moe_lora_format,
    )

    generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for vllm.model_executor.layers.pooler.activations."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.model_executor.layers.pooler.activations import (
+    LambdaPoolerActivation,
+    PoolerClassify,
+    PoolerIdentity,
+    PoolerMultiLabelClassify,
+    PoolerNormalize,
+    get_act_fn,
+    resolve_classifier_act_fn,
+)
+
+
+# ---------------------------------------------------------------------------
+# PoolerIdentity
+# ---------------------------------------------------------------------------
+class TestPoolerIdentity:
+    def test_returns_input_unchanged(self):
+        pooler = PoolerIdentity()
+        x = torch.randn(4, 128)
+        out = pooler(x)
+        assert torch.equal(out, x)
+
+    def test_forward_list(self):
+        pooler = PoolerIdentity()
+        tensors = [torch.randn(128), torch.randn(256)]
+        out = pooler(tensors)
+        assert len(out) == 2
+        for orig, result in zip(tensors, out):
+            assert torch.equal(orig, result)
+
+
+# ---------------------------------------------------------------------------
+# PoolerNormalize
+# ---------------------------------------------------------------------------
+class TestPoolerNormalize:
+    def test_output_has_unit_norm(self):
+        pooler = PoolerNormalize()
+        x = torch.randn(4, 128)
+        out = pooler(x)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(4), atol=1e-5)
+
+    def test_single_vector(self):
+        pooler = PoolerNormalize()
+        x = torch.randn(1, 64)
+        out = pooler(x)
+        norm = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norm, torch.ones(1), atol=1e-5)
+
+    def test_forward_list(self):
+        pooler = PoolerNormalize()
+        tensors = [torch.randn(1, 64), torch.randn(1, 128)]
+        out = pooler(tensors)
+        for t in out:
+            norm = torch.linalg.norm(t, dim=-1)
+            assert torch.allclose(norm, torch.ones(1), atol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# PoolerMultiLabelClassify
+# ---------------------------------------------------------------------------
+class TestPoolerMultiLabelClassify:
+    def test_output_in_zero_one(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.randn(4, 10)
+        out = pooler(x)
+        assert (out >= 0).all() and (out <= 1).all()
+
+    def test_large_positive_maps_near_one(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.full((1, 3), 100.0)
+        out = pooler(x)
+        assert torch.allclose(out, torch.ones(1, 3), atol=1e-4)
+
+    def test_large_negative_maps_near_zero(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.full((1, 3), -100.0)
+        out = pooler(x)
+        assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# PoolerClassify
+# ---------------------------------------------------------------------------
+class TestPoolerClassify:
+    def test_infers_from_shape_when_num_labels_none(self):
+        pooler = PoolerClassify(num_labels=None)
+        assert pooler.num_labels is None
+        x = torch.randn(2, 5)
+        out = pooler(x)
+        sums = out.sum(dim=-1)
+        assert torch.allclose(sums, torch.ones(2), atol=1e-5)
+
+    def test_sigmoid_when_num_labels_lt_2(self):
+        pooler = PoolerClassify(num_labels=1)
+        x = torch.zeros(1, 1)
+        out = pooler(x)
+        assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5)
+
+    def test_num_labels_zero_uses_sigmoid(self):
+        pooler = PoolerClassify(num_labels=0)
+        assert pooler.num_labels == 0
+        x = torch.zeros(1, 3)
+        out = pooler(x)
+        assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5)
+
+    def test_num_labels_ge_2_uses_softmax(self):
+        pooler = PoolerClassify(num_labels=4)
+        assert pooler.num_labels == 4
+        x = torch.randn(2, 4)
+        out = pooler(x)
+        sums = out.sum(dim=-1)
+        assert torch.allclose(sums, torch.ones(2), atol=1e-5)
+
+    def test_default_num_labels_is_none(self):
+        pooler = PoolerClassify()
+        assert pooler.num_labels is None
+
+
+# ---------------------------------------------------------------------------
+# LambdaPoolerActivation
+# ---------------------------------------------------------------------------
+class TestLambdaPoolerActivation:
+    def test_applies_custom_fn(self):
+        pooler = LambdaPoolerActivation(nn.ReLU())
+        x = torch.tensor([[-1.0, 2.0, -3.0]])
+        out = pooler(x)
+        expected = torch.tensor([[0.0, 2.0, 0.0]])
+        assert torch.equal(out, expected)
+
+    def test_forward_list(self):
+        pooler = LambdaPoolerActivation(nn.ReLU())
+        tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])]
+        out = pooler(tensors)
+        assert torch.equal(out[0], torch.tensor([0.0, 2.0]))
+        assert torch.equal(out[1], torch.tensor([3.0, 0.0]))
+
+
+# ---------------------------------------------------------------------------
+# get_act_fn factory
+# ---------------------------------------------------------------------------
+class TestGetActFn:
+    @staticmethod
+    def _make_config(**kwargs):
+        return SimpleNamespace(**kwargs)
+
+    def test_regression(self):
+        cfg = self._make_config(problem_type="regression")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerIdentity)
+
+    def test_single_label_classification(self):
+        cfg = self._make_config(
+            problem_type="single_label_classification", num_labels=3
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+        assert result.num_labels == 3
+
+    def test_multi_label_classification(self):
+        cfg = self._make_config(problem_type="multi_label_classification")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerMultiLabelClassify)
+
+    def test_sentence_transformers_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={
+                "activation_fn": "torch.nn.modules.activation.Sigmoid"
+            },
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_sbert_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sbert_ce_default_activation_function=(
+                "torch.nn.modules.activation.Sigmoid"
+            ),
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_default_fallback(self):
+        cfg = self._make_config(problem_type="")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_sentence_transformers_takes_priority(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"},
+            sbert_ce_default_activation_function=(
+                "torch.nn.modules.activation.Sigmoid"
+            ),
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerIdentity)
+
+    def test_rejects_non_torch_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={"activation_fn": "os.system"},
+        )
+        with pytest.raises(AssertionError, match="restricted"):
+            get_act_fn(cfg)
+
+
+# ---------------------------------------------------------------------------
+# resolve_classifier_act_fn
+# ---------------------------------------------------------------------------
+class TestResolveClassifierActFn:
+    def test_delegates_to_get_act_fn_when_none(self):
+        model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(num_labels=3, problem_type="")
+        )
+        result = resolve_classifier_act_fn(model_config, act_fn=None)
+        assert isinstance(result, PoolerClassify)
+        assert result.num_labels == 3
+
+    def test_passes_through_provided_act_fn(self):
+        custom = PoolerIdentity()
+        result = resolve_classifier_act_fn(None, act_fn=custom)
+        assert result is custom
@@ -6,8 +6,15 @@ from weakref import WeakKeyDictionary, ref

 import pytest
 import torch
+from torch.nn.parameter import UninitializedParameter

+import vllm.model_executor.model_loader.reload.meta as reload_meta
 from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.model_loader.reload.layerwise import (
+    finalize_layerwise_reload,
+    initialize_layerwise_reload,
+    record_metadata_for_reloading,
+)
 from vllm.model_executor.model_loader.reload.meta import (
    capture_layer_to_meta,
    get_numel_loaded,
@@ -21,6 +28,38 @@ from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
 from vllm.platforms import current_platform


+class _AliasedBufferLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        weight = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        self.weight = torch.nn.Parameter(weight)
+        self.register_buffer(
+            "weight_view", self.weight.detach().view(-1), persistent=False
+        )
+
+
+class _ParentAliasedChildBufferLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.ones(1))
+        self.conv1d = torch.nn.Linear(3, 2, bias=False)
+        self.conv1d.weight.data.copy_(
+            torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        )
+        self.register_buffer(
+            "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False
+        )
+
+
+class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer):
+    def __init__(self):
+        super().__init__()
+        self.child = torch.nn.Module()
+        self.child.register_parameter(
+            "lazy_weight", UninitializedParameter(requires_grad=False)
+        )
+
+
 def test_move_metatensors():
    tensor = torch.empty((1, 2, 3))
    meta_tensor = to_meta_tensor(tensor)
@@ -123,6 +162,81 @@ def test_get_numel_loaded():
    assert ret == "value"


+def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch):
+    layer = _AliasedBufferLayer()
+    model = torch.nn.Sequential(layer)
+    loaded_weight = torch.full_like(layer.weight, 7.0)
+
+    def materialize_with_sentinel(meta_tensor):
+        tensor = torch.empty_strided(
+            size=tuple(meta_tensor.size()),
+            stride=tuple(meta_tensor.stride()),
+            dtype=meta_tensor.dtype,
+            requires_grad=False,
+        )
+        tensor.fill_(-123.0)
+        tensor.__class__ = meta_tensor.__class__
+        tensor.__dict__ = meta_tensor.__dict__.copy()
+        return tensor
+
+    monkeypatch.setattr(
+        reload_meta, "materialize_meta_tensor", materialize_with_sentinel
+    )
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    layer.weight.weight_loader(layer.weight, loaded_weight)
+    finalize_layerwise_reload(model, model_config=None)
+
+    assert torch.equal(layer.weight, loaded_weight)
+    assert layer.weight_view.untyped_storage().data_ptr() == (
+        layer.weight.untyped_storage().data_ptr()
+    )
+
+
+def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs():
+    layer = _AliasedBufferWithUninitializedChildLayer()
+
+    _, buffers = capture_layer_to_meta(layer)
+
+    assert "weight_view" not in buffers
+
+
+def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch):
+    layer = _ParentAliasedChildBufferLayer()
+    model = torch.nn.Sequential(layer)
+    loaded_conv = torch.full_like(layer.conv1d.weight, 7.0)
+    loaded_scale = torch.full_like(layer.scale, 3.0)
+
+    def materialize_with_sentinel(meta_tensor):
+        tensor = torch.empty_strided(
+            size=tuple(meta_tensor.size()),
+            stride=tuple(meta_tensor.stride()),
+            dtype=meta_tensor.dtype,
+            requires_grad=False,
+        )
+        tensor.fill_(-123.0)
+        tensor.__class__ = meta_tensor.__class__
+        tensor.__dict__ = meta_tensor.__dict__.copy()
+        return tensor
+
+    monkeypatch.setattr(
+        reload_meta, "materialize_meta_tensor", materialize_with_sentinel
+    )
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv)
+    layer.scale.weight_loader(layer.scale, loaded_scale)
+    finalize_layerwise_reload(model, model_config=None)
+
+    assert torch.equal(layer.conv1d.weight, loaded_conv)
+    assert torch.equal(layer.conv_weights, loaded_conv.view(-1))
+    assert layer.conv_weights.untyped_storage().data_ptr() == (
+        layer.conv1d.weight.untyped_storage().data_ptr()
+    )
+
+
@pytest.mark.parametrize(
    "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
 )
@@ -7,7 +7,7 @@ from scipy.spatial.distance import cosine
 from vllm import LLM, SamplingParams
 from vllm.config import ModelConfig

-from ....utils import RemoteOpenAIServer
+from ....utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
 from .embed_utils import run_client_embeddings

 MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
@@ -126,9 +126,15 @@ def test_gritlm_offline_embedding(vllm_runner):
 async def test_gritlm_api_server_embedding():
    queries, q_instruction, documents, d_instruction = get_test_data()

-    args = ["--runner", "pooling", "--max_model_len", str(MAX_MODEL_LEN)]
+    args = [
+        "--runner",
+        "pooling",
+        "--max_model_len",
+        str(MAX_MODEL_LEN),
+        *ROCM_EXTRA_ARGS,
+    ]

-    with RemoteOpenAIServer(MODEL_NAME, args) as server:
+    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=ROCM_ENV_OVERRIDES) as server:
        client_embedding = server.get_async_client()

        d_rep = await run_client_embeddings(
@@ -17,6 +17,7 @@ from vllm.entrypoints.pooling.scoring.protocol import RerankResponse
 os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"

 TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template")
+ExpectedPromptTokens = int | tuple[int, ...]

 long_query = "What is the capital of France?" * 20
 long_doc = "The capital of France is Paris. " * 20
@@ -26,10 +27,10 @@ long_doc = "The capital of France is Paris. " * 20
 class TestConfig:
    model: str
    args: list[str]
-    without_truncated_prompt_tokens: int
-    with_max_tokens_per_query_prompt_tokens: int
-    with_max_tokens_per_doc_prompt_tokens: int
-    with_max_tokens_per_query_and_doc_prompt_tokens: int
+    without_truncated_prompt_tokens: ExpectedPromptTokens
+    with_max_tokens_per_query_prompt_tokens: ExpectedPromptTokens
+    with_max_tokens_per_doc_prompt_tokens: ExpectedPromptTokens
+    with_max_tokens_per_query_and_doc_prompt_tokens: ExpectedPromptTokens


 RERANK_CONFIGS = [
@@ -79,8 +80,10 @@ RERANK_CONFIGS = [
            "512",
            "--trust-remote-code",
        ],
-        without_truncated_prompt_tokens=286,
-        with_max_tokens_per_query_prompt_tokens=156,
+        # This model has produced both prompt-token totals in CI/local cache;
+        # keep truncation checks exact while tolerating the boundary delta.
+        without_truncated_prompt_tokens=(285, 286),
+        with_max_tokens_per_query_prompt_tokens=(155, 156),
        with_max_tokens_per_doc_prompt_tokens=155,
        with_max_tokens_per_query_and_doc_prompt_tokens=25,
    ),
@@ -115,6 +118,13 @@ RERANK_CONFIGS = [
 ]


+def assert_prompt_tokens(actual: int, expected: ExpectedPromptTokens) -> None:
+    if isinstance(expected, int):
+        assert actual == expected
+    else:
+        assert actual in expected
+
+
@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model)
 def server(request):
    config: TestConfig = request.param
@@ -136,7 +146,10 @@ def test_without_truncated(server):
    assert rerank.id is not None
    assert rerank.results is not None
    assert len(rerank.results) == 1
-    assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens
+    assert_prompt_tokens(
+        rerank.usage.prompt_tokens,
+        config.without_truncated_prompt_tokens,
+    )


 def test_max_tokens_per_query(server):
@@ -158,7 +171,10 @@ def test_max_tokens_per_query(server):
    assert rerank.id is not None
    assert rerank.results is not None
    assert len(rerank.results) == 1
-    assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens
+    assert_prompt_tokens(
+        rerank.usage.prompt_tokens,
+        config.with_max_tokens_per_query_prompt_tokens,
+    )


 def test_max_tokens_per_doc(server):
@@ -180,7 +196,10 @@ def test_max_tokens_per_doc(server):
    assert rerank.id is not None
    assert rerank.results is not None
    assert len(rerank.results) == 1
-    assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens
+    assert_prompt_tokens(
+        rerank.usage.prompt_tokens,
+        config.with_max_tokens_per_doc_prompt_tokens,
+    )


 def test_max_tokens_per_query_and_doc(server):
@@ -203,7 +222,7 @@ def test_max_tokens_per_query_and_doc(server):
    assert rerank.id is not None
    assert rerank.results is not None
    assert len(rerank.results) == 1
-    assert (
-        rerank.usage.prompt_tokens
-        == config.with_max_tokens_per_query_and_doc_prompt_tokens
+    assert_prompt_tokens(
+        rerank.usage.prompt_tokens,
+        config.with_max_tokens_per_query_and_doc_prompt_tokens,
    )
@@ -34,6 +34,7 @@ WINDOW_ATTN_IMAGE_PROMPT = qwen2_5_vl_chat_template(
    IMAGE_PLACEHOLDER,
    "Describe the image.",
 )
+IMAGE_ONLY_LIMIT_MM_PER_PROMPT = {"image": 1, "video": 0}


 def _window_attention_regression_image():
@@ -193,7 +194,7 @@ def test_qwen2_5_vl_window_attention_image(
        runner="generate",
        max_model_len=4096,
        dtype=dtype,
-        limit_mm_per_prompt={"image": 1},
+        limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
        compilation_config=_encoder_cudagraph_config(max_vision_items=1),
    ) as vllm_model:
        outputs = vllm_model.generate_greedy(prompt, max_tokens, images=images)
@@ -231,7 +232,7 @@ def test_qwen2_5_vl_window_attention_image_batch(
        max_model_len=4096,
        max_num_seqs=2,
        dtype=dtype,
-        limit_mm_per_prompt={"image": 1},
+        limit_mm_per_prompt=IMAGE_ONLY_LIMIT_MM_PER_PROMPT,
        compilation_config=_encoder_cudagraph_config(max_vision_items=2),
    ) as vllm_model:
        outputs = vllm_model.generate_greedy(prompts, max_tokens, images=images)
@@ -41,6 +41,13 @@ def qwen_vl_chat_template(content: str) -> str:
    return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"


+def step3_vl_chat_template(content: str) -> str:
+    return (
+        "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
+        f"<im_patch>{content} <|EOT|><|BOT|>assistant\n"
+    )
+
+
 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
    "qwen2_5_vl": VitCudagraphTestConfig(
        model="Qwen/Qwen2.5-VL-3B-Instruct",
@@ -90,6 +97,11 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
        needs_video_metadata=False,
        marks=[pytest.mark.core_model],
    ),
+    "step3_vl": VitCudagraphTestConfig(
+        model="stepfun-ai/Step3-VL-10B",
+        image_prompt=step3_vl_chat_template("What is in this image?"),
+        video_prompt=None,
+    ),
 }


@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
 import pytest

 from tests.quantization.utils import is_quant_method_supported
-from vllm.v1.attention.backends.fa_utils import flash_attn_supports_fp8
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from ..utils import check_logprobs_close


@@ -68,7 +68,13 @@ def test_models(
    if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
        pytest.skip(f"{kv_cache_dtype} is currently not supported on ROCm/HIP.")

-    if not flash_attn_supports_fp8():
+    if not (
+        current_platform.is_xpu()
+        or (
+            get_flash_attn_version() == 3
+            and current_platform.is_device_capability_family(90)
+        )
+    ):
        pytest.skip(
            f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
        )
@@ -90,9 +90,9 @@ def test_models(example_prompts, model_name) -> None:
 EAGER = [True, False]

 SM_100_NVFP4_BACKENDS = [
-    "flashinfer-cudnn",
-    "flashinfer-trtllm",
-    "flashinfer-cutlass",
+    "flashinfer_cudnn",
+    "flashinfer_trtllm",
+    "flashinfer_cutlass",
 ]


@@ -102,12 +102,12 @@ SM_100_NVFP4_BACKENDS = [
    "backend",
    [
        "emulation",
-        "flashinfer-cudnn",
-        "flashinfer-trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
-        "flashinfer-cutlass",
+        "flashinfer_cudnn",
+        "flashinfer_trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
+        "flashinfer_cutlass",
    ],
 )
-def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
+def test_nvfp4(vllm_runner, model, eager, backend):
    if (
        not current_platform.has_device_capability(100)
        and backend in SM_100_NVFP4_BACKENDS
@@ -116,8 +116,7 @@ def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
            f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False"
        )

-    monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)
-    with vllm_runner(model, enforce_eager=eager) as llm:
+    with vllm_runner(model, enforce_eager=eager, linear_backend=backend) as llm:
        output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
    assert output[0][1] == "1 2 3 4 5 6"

@@ -122,14 +122,17 @@ class TestGGUFModelLoader:

        mock_hf_download.return_value = "/downloaded/model.gguf"

-        # Create a simple mock ModelConfig with only the model attribute
        model_config = MagicMock()
        model_config.model = "unsloth/Qwen3-0.6B-GGUF/model.gguf"
+        model_config.revision = "abc123"

        result = loader._prepare_weights(model_config)
        assert result == "/downloaded/model.gguf"
        mock_hf_download.assert_called_once_with(
-            repo_id="unsloth/Qwen3-0.6B-GGUF", filename="model.gguf"
+            repo_id="unsloth/Qwen3-0.6B-GGUF",
+            filename="model.gguf",
+            revision="abc123",
+            cache_dir=None,
        )

    @patch("vllm.config.model.get_hf_image_processor_config", return_value=None)
@@ -15,6 +15,7 @@ MODELS = [
    "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
    "Qwen/Qwen3-0.6B-FP8",  # FP8 W8A16 block-quantized linear
    "Qwen/Qwen3-30B-A3B-FP8",  # FP8 W8A16 block-quantized MoE
+    "openai/gpt-oss-20b",  # MXFP4 W4A16
 ]
 DTYPE = ["bfloat16"]

@@ -1,6 +1,90 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950).
+
+These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    RoutingMethodType,
+)
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    Mxfp4MoeBackend,
+    select_mxfp4_moe_backend,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kMxfp4Dynamic,
+)
+from vllm.platforms import current_platform
+
+ROCM_AVAILABLE = current_platform.is_rocm()
+ROCM_GFX950 = False
+ROCM_AITER_AVAILABLE = False
+
+if ROCM_AVAILABLE:
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.platforms.rocm import on_gfx950
+
+    ROCM_GFX950 = on_gfx950()
+    ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()


-def test_mi355_moe():
-    print("TODO: add tests for Mi355 MoE quantization")
+def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+    return FusedMoEConfig(
+        num_experts=8,
+        experts_per_token=2,
+        hidden_dim=256,
+        intermediate_size_per_partition=256,
+        num_local_experts=8,
+        num_logical_experts=8,
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        activation=MoEActivation.SILU,
+        in_dtype=torch.bfloat16,
+        device="cuda",
+        routing_method=RoutingMethodType.Renormalize,
+        moe_backend=moe_backend,
+    )
+
+
+@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
+@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")
+def test_w4a4_dispatches_to_aiter():
+    """With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""
+    config = _make_w4a4_moe_config()
+    backend, experts_cls = select_mxfp4_moe_backend(
+        config, activation_key=kMxfp4Dynamic
+    )
+    assert backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4
+    assert experts_cls is not None
+
+
+@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
+@pytest.mark.skipif(
+    ROCM_AITER_AVAILABLE,
+    reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)",
+)
+def test_w4a4_raises_without_aiter_and_no_moe_backend():
+    """Without AITER and no --moe-backend, raises NotImplementedError
+    with hint to use --moe-backend emulation."""
+    config = _make_w4a4_moe_config()
+    with pytest.raises(NotImplementedError, match="--moe-backend emulation"):
+        select_mxfp4_moe_backend(config, activation_key=kMxfp4Dynamic)
+
+
+@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
+def test_w4a4_dispatches_to_emulation_with_moe_backend():
+    """With --moe-backend emulation, W4A4 selects EMULATION."""
+    config = _make_w4a4_moe_config(moe_backend="emulation")
+    backend, experts_cls = select_mxfp4_moe_backend(
+        config, activation_key=kMxfp4Dynamic
+    )
+    assert backend == Mxfp4MoeBackend.EMULATION
+    assert experts_cls is not None
@@ -148,6 +148,13 @@ def test_is_default_v2_model_runner_model(model_config, expected):
    assert VllmConfig._is_default_v2_model_runner_model(config) is expected


+def test_use_v2_model_runner_defaults_to_v1_when_kv_connector_present():
+    config = SimpleNamespace(kv_transfer_config=object())
+    with patch.object(envs, "VLLM_USE_V2_MODEL_RUNNER", None):
+        result = VllmConfig.use_v2_model_runner.fget(config)
+    assert result is False
+
+
@pytest.mark.skip_global_cleanup
 def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping(
    monkeypatch,
@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.tool_parsers.apertus_tool_parser import (
+    TOOL_CALLS_PREFIX,
+    TOOL_CALLS_SUFFIX,
+    ApertusToolParser,
+)
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_tokenizer():
+    tokenizer = MagicMock()
+    tokenizer.encode.return_value = [1, 2, 3]
+    # Include the tool call tokens in the vocab for the parser
+    tokenizer.get_vocab.return_value = {TOOL_CALLS_PREFIX: 100, TOOL_CALLS_SUFFIX: 101}
+    return tokenizer
+
+
+@pytest.fixture
+def parser(mock_tokenizer):
+    return ApertusToolParser(mock_tokenizer)
+
+
+@pytest.fixture
+def mock_request():
+    request = MagicMock(spec=ChatCompletionRequest)
+    request.tools = []
+    request.tool_choice = "auto"
+    return request
+
+
+# ---------------------------------------------------------------------------
+# Non-streaming extraction tests
+# ---------------------------------------------------------------------------
+
+
+class TestExtractToolCalls:
+    def test_no_tool_calls(self, parser, mock_request):
+        model_output = "Hello, how can I help you today?"
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is False
+        assert result.tool_calls == []
+        assert result.content == model_output
+
+    def test_single_tool_call(self, parser, mock_request):
+        model_output = (
+            '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "London"}
+
+    def test_multiple_arguments(self, parser, mock_request):
+        model_output = (
+            '<|tools_prefix|>[{"get_weather": '
+            '{"location": "San Francisco", '
+            '"unit": "celsius"}}]<|tools_suffix|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "San Francisco", "unit": "celsius"}
+
+    def test_text_before_tool_call(self, parser, mock_request):
+        model_output = (
+            "Let me check the weather for you. "
+            '<|tools_prefix|>[{"get_weather": {"location": "Paris"}}]<|tools_suffix|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert result.content == "Let me check the weather for you."
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+
+    def test_multiple_tool_calls(self, parser, mock_request):
+        model_output = (
+            '<|tools_prefix|>[{"get_weather": '
+            '{"location": "London"}}, '
+            '{"get_time": {"location": "London"}}]<|tools_suffix|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 2
+        assert result.tool_calls[0].function.name == "get_weather"
+        assert result.tool_calls[1].function.name == "get_time"
+
+    def test_nested_arguments(self, parser, mock_request):
+        model_output = (
+            '<|tools_prefix|>[{"complex_function": '
+            '{"nested": {"inner": "value"}, '
+            '"list": ["a", "b"]}}]<|tools_suffix|>'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "complex_function"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]}
+
+    def test_incomplete_tool_call(self, parser, mock_request):
+        model_output = '<|tools_prefix|>[{"get_weather": {"location": "London"}'
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "London"}
+
+    def test_missing_tool_suffix(self, parser, mock_request):
+        model_output = (
+            '<|tools_prefix|>[{"get_weather": '
+            '{"location": "San Francisco", "unit": "celsius"}}]'
+        )
+        result = parser.extract_tool_calls(model_output, mock_request)
+
+        assert result.tools_called is True
+        assert len(result.tool_calls) == 1
+        assert result.tool_calls[0].function.name == "get_weather"
+        args = json.loads(result.tool_calls[0].function.arguments)
+        assert args == {"location": "San Francisco", "unit": "celsius"}
+
+
+# ---------------------------------------------------------------------------
+# Streaming extraction tests
+# ---------------------------------------------------------------------------
+
+
+class TestStreamingExtraction:
+    def _simulate_streaming(
+        self, parser: ApertusToolParser, mock_request: Any, chunks: list[str]
+    ) -> list[tuple[Any, str]]:
+        results: list[tuple[Any, str]] = []
+        previous_text: str = ""
+        previous_token_ids: list[int] = []
+
+        for chunk in chunks:
+            current_text = previous_text + chunk
+            # Simulate a token ID sequence matching the chunk progression
+            delta_token_ids: list[int] = [0] * max(1, len(chunk) // 4)
+            current_token_ids = previous_token_ids + delta_token_ids
+
+            delta = parser.extract_tool_calls_streaming(
+                previous_text=previous_text,
+                current_text=current_text,
+                delta_text=chunk,
+                previous_token_ids=tuple(previous_token_ids),
+                current_token_ids=tuple(current_token_ids),
+                delta_token_ids=tuple(delta_token_ids),
+                request=mock_request,
+            )
+            results.append((delta, current_text))
+            previous_text = current_text
+            previous_token_ids = list(current_token_ids)
+
+        return results
+
+    def _collect_tool_calls(self, results) -> dict[int, dict[str, Any]]:
+        """Properly tracks and concatenates streamed tool arguments by their Index."""
+        tool_calls = {}
+        for delta, _ in results:
+            if not delta or not getattr(delta, "tool_calls", None):
+                continue
+
+            for tc in delta.tool_calls:
+                idx = (
+                    tc.get("index", 0)
+                    if isinstance(tc, dict)
+                    else getattr(tc, "index", 0)
+                )
+                func = (
+                    tc.get("function", {})
+                    if isinstance(tc, dict)
+                    else getattr(tc, "function", None)
+                )
+                if not func:
+                    continue
+
+                name = (
+                    func.get("name")
+                    if isinstance(func, dict)
+                    else getattr(func, "name", None)
+                )
+                args = (
+                    func.get("arguments")
+                    if isinstance(func, dict)
+                    else getattr(func, "arguments", None)
+                )
+
+                if idx not in tool_calls:
+                    tool_calls[idx] = {"name": "", "arguments": ""}
+
+                if name:
+                    tool_calls[idx]["name"] += name
+                if args:
+                    tool_calls[idx]["arguments"] += args
+
+        return tool_calls
+
+    def _collect_content(self, results) -> str:
+        """Collects generated normal text outside of the tool calls."""
+        return "".join(
+            delta.content
+            for delta, _ in results
+            if delta and getattr(delta, "content", None)
+        )
+
+    def test_basic_streaming_single_tool(self, parser, mock_request):
+        chunks = [
+            "<|tools_prefix|>",
+            '[{"get_weather": ',
+            '{"location": "Paris, ',
+            'France"}}]',
+            "<|tools_suffix|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        tcs = self._collect_tool_calls(results)
+
+        assert len(tcs) == 1
+        assert tcs[0]["name"] == "get_weather"
+        assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"}
+
+    def test_streaming_missing_tool_suffix(self, parser, mock_request):
+        chunks = [
+            "<|tools_prefix|>",
+            '[{"get_weather": ',
+            '{"location": "Paris, ',
+            'France"}}]',
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        tcs = self._collect_tool_calls(results)
+
+        assert len(tcs) == 1
+        assert tcs[0]["name"] == "get_weather"
+        assert json.loads(tcs[0]["arguments"]) == {"location": "Paris, France"}
+
+    def test_streaming_partial_tag_buffering_missing_tool_suffix(
+        self, parser, mock_request
+    ):
+        chunks = ["Content", "<|tools_", "prefix|>", '[{"f": ', '{"a": 1}}]']
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        content = self._collect_content(results)
+
+        assert "Content" in content
+        assert "<|tools_prefix|>" not in content
+        assert "<|tools_suffix|>" not in content
+
+        tcs = self._collect_tool_calls(results)
+
+        assert len(tcs) == 1
+        assert tcs[0]["name"] == "f"
+        assert json.loads(tcs[0]["arguments"]) == {"a": 1}
+
+    def test_streaming_multi_tool(self, parser, mock_request):
+        chunks = [
+            "<|tools_prefix|>",
+            '[{"get_weather": {"location": "Tokyo"}}',
+            ', {"get_time": {"location": "Tokyo"}}]',
+            "<|tools_suffix|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        tcs = self._collect_tool_calls(results)
+
+        assert len(tcs) == 2
+        assert tcs[0]["name"] == "get_weather"
+        assert json.loads(tcs[0]["arguments"]) == {"location": "Tokyo"}
+        assert tcs[1]["name"] == "get_time"
+        assert json.loads(tcs[1]["arguments"]) == {"location": "Tokyo"}
+
+    def test_streaming_text_before_tool_call(self, parser, mock_request):
+        chunks = [
+            "Let me check ",
+            "the weather. ",
+            "<|tools_prefix|>",
+            '[{"get_weather": {"location": "London"}}]',
+            "<|tools_suffix|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        content = self._collect_content(results)
+
+        assert content.strip() == "Let me check the weather."
+        tcs = self._collect_tool_calls(results)
+
+        assert len(tcs) == 1
+        assert tcs[0]["name"] == "get_weather"
+        assert json.loads(tcs[0]["arguments"]) == {"location": "London"}
+
+    def test_streaming_partial_tag_buffering(self, parser, mock_request):
+        chunks = [
+            "Content",
+            "<|tools_",
+            "prefix|>",
+            '[{"f": {"a": 1}}]',
+            "<|tools_suf",
+            "fix|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+        content = self._collect_content(results)
+
+        assert "Content" in content
+        assert "<|tools_prefix|>" not in content
+        assert "<|tools_suffix|>" not in content
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 1
+        assert tc[0]["name"] == "f"
+        assert json.loads(tc[0]["arguments"]) == {"a": 1}
+
+    # ---------------------------------------------------------------------------
+    # Edge Cases: Multi-Token Prediction (MTP) & vLLM Chunking Anomalies
+    # ---------------------------------------------------------------------------
+
+    def test_mtp_streaming_massive_chunk(self, parser, mock_request):
+        """Simulates MTP predicting text, tool calls,
+        and trailing text all in a single chunk."""
+        chunks = [
+            "Sure! "
+            '<|tools_prefix|>[{"get_weather": {"location": "London"}}]<|tools_suffix|>'
+        ]
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        content = self._collect_content(results)
+        assert "Sure! " in content
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 1
+        assert tc[0]["name"] == "get_weather"
+        assert json.loads(tc[0]["arguments"]) == {"location": "London"}
+
+    def test_mtp_streaming_multiple_tools_burst(self, parser, mock_request):
+        """Simulates MTP predicting an array of multiple tools in one single chunk."""
+        chunks = [
+            '<|tools_prefix|>[{"get_weather": '
+            '{"location": "London"}}, '
+            '{"get_time": {"location": "Paris"}}]<|tools_suffix|>'
+        ]
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 2
+        assert tc[0]["name"] == "get_weather"
+        assert json.loads(tc[0]["arguments"]) == {"location": "London"}
+        assert tc[1]["name"] == "get_time"
+        assert json.loads(tc[1]["arguments"]) == {"location": "Paris"}
+
+    def test_mtp_streaming_skip_and_catch_up(self, parser, mock_request):
+        """Simulates MTP chunks that jump over entire tools
+        (e.g., from middle of tool 1 to middle of tool 3)."""
+        chunks = [
+            '<|tools_prefix|>[{"t1": {"a": 1}',
+            '}, {"t2": {"b": 2}}, {"t3": {"c": 3',
+            "}}]<|tools_suffix|>",
+        ]
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 3
+        assert tc[0]["name"] == "t1"
+        assert json.loads(tc[0]["arguments"]) == {"a": 1}
+        assert tc[1]["name"] == "t2"
+        assert json.loads(tc[1]["arguments"]) == {"b": 2}
+        assert tc[2]["name"] == "t3"
+        assert json.loads(tc[2]["arguments"]) == {"c": 3}
+
+    def test_vllm_streaming_character_by_character(self, parser, mock_request):
+        """Simulates worst-case vLLM fragmentation where
+        chunks arrive character-by-character."""
+        text = (
+            'Hi <|tools_prefix|>[{"get_weather": '
+            '{"location": "London"}}]<|tools_suffix|> '
+        )
+        chunks = list(text)
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        content = self._collect_content(results)
+        assert "Hi" in content
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 1
+        assert tc[0]["name"] == "get_weather"
+        assert json.loads(tc[0]["arguments"]) == {"location": "London"}
+
+    def test_vllm_streaming_empty_deltas(self, parser, mock_request):
+        """Simulates vLLM stream producing empty string chunks
+        (e.g., hidden tokens or artifacts)."""
+        chunks = [
+            "Wait",
+            "",
+            "<|tools_prefix|>",
+            "",
+            '[{"get_weather": ',
+            "",
+            '{"location": "London"}}]',
+            "<|tools_suffix|>",
+        ]
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        content = self._collect_content(results)
+        assert content == "Wait"
+
+        tc = self._collect_tool_calls(results)
+        assert len(tc) == 1
+        assert tc[0]["name"] == "get_weather"
+        assert json.loads(tc[0]["arguments"]) == {"location": "London"}
@@ -16,39 +16,60 @@ from types import SimpleNamespace

 import torch

+from tests.v1.attention.utils import MockMambaBuilder
 from vllm.config.compilation import CUDAGraphMode
-from vllm.v1.attention.backends.mamba_attn import (
-    BaseMambaAttentionMetadata,
-    BaseMambaAttentionMetadataBuilder,
-)
+from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata
 from vllm.v1.kv_cache_interface import MambaSpec


-class _ConcreteMambaBuilder(
-    BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]
+def _make_vllm_config(
+    max_model_len: int,
+    max_num_seqs: int,
+    num_speculative_tokens: int = 0,
+    block_size: int | None = None,
 ):
-    """Minimal concrete subclass for testing (base class is ABC)."""
-
-    metadata_cls = BaseMambaAttentionMetadata
-
-
-def _make_vllm_config(block_size, max_model_len, max_num_seqs):
    """Create a minimal mock VllmConfig with only the fields the builder
    accesses, avoiding any model download / HF config inspection."""
+    speculative_config = (
+        SimpleNamespace(
+            num_speculative_tokens=num_speculative_tokens,
+            parallel_drafting=False,
+        )
+        if num_speculative_tokens > 0
+        else None
+    )
    return SimpleNamespace(
-        cache_config=SimpleNamespace(mamba_cache_mode="all"),
+        cache_config=SimpleNamespace(
+            block_size=block_size,
+            mamba_cache_mode="all",
+        ),
        compilation_config=SimpleNamespace(
            cudagraph_mode=CUDAGraphMode.FULL,
            max_cudagraph_capture_size=None,
        ),
-        speculative_config=None,
-        num_speculative_tokens=0,
+        speculative_config=speculative_config,
+        num_speculative_tokens=num_speculative_tokens,
        parallel_config=SimpleNamespace(decode_context_parallel_size=1),
        scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
        model_config=SimpleNamespace(max_model_len=max_model_len),
    )


+def test_mamba_single_token_prompt_runs_as_prefill():
+    seq_lens = [8, 9, 1]
+    config = _make_vllm_config(256, len(seq_lens), block_size=16)
+    metadata = MockMambaBuilder.build_mamba_metadata(
+        config,
+        seq_lens=seq_lens,
+        query_lens=[1] * len(seq_lens),
+        is_prefilling=[False, False, True],
+    )
+
+    assert metadata.num_decodes == 2
+    assert metadata.num_prefills == 1
+    assert metadata.has_initial_states_p.tolist() == [False]
+
+
 def test_update_block_table_copies_block_idx_to_persistent_buffers():
    """update_block_table() must write block_idx tensors to the current
    builder's persistent buffers, not leave them pointing to a different
@@ -59,7 +80,7 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
    num_reqs = 4
    device = torch.device("cpu")

-    vllm_config = _make_vllm_config(block_size, max_model_len, num_reqs)
+    vllm_config = _make_vllm_config(max_model_len, num_reqs)

    spec = MambaSpec(
        block_size=block_size,
@@ -69,8 +90,8 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
    )

    # Two builders simulating two KV cache groups with the same MambaSpec.
-    builder_a = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
-    builder_b = _ConcreteMambaBuilder(spec, ["layer1"], vllm_config, device)
+    builder_a = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+    builder_b = MockMambaBuilder(spec, ["layer1"], vllm_config, device)

    # Sanity: each builder has its own persistent buffer.
    assert (
@@ -106,6 +127,7 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
        block_idx_last_computed_token=(
            builder_a.block_idx_last_computed_token[:num_reqs]
        ),
+        block_idx_last_scheduled_token_prev_step=None,
        seq_lens=seq_lens,
    )

@@ -149,3 +171,261 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
        metadata_b.block_idx_last_computed_token,
        block_idx_vals,
    )
+
+
+def test_state_indices_tensor_d_includes_num_speculative_blocks():
+    """Regression test for https://github.com/vllm-project/vllm/issues/39809
+    bug 1: with mamba_cache_mode='all' and speculative decoding enabled,
+    the cudagraph buffer for state_indices_tensor_d must allocate the same
+    per-request column count as the runtime block table, which includes
+    num_speculative_blocks trailing scratch columns."""
+
+    block_size = 16
+    max_model_len = 256
+    max_num_seqs = 4
+    num_speculative_tokens = 1
+    num_speculative_blocks = 2
+    device = torch.device("cpu")
+
+    vllm_config = _make_vllm_config(
+        max_model_len,
+        max_num_seqs,
+        num_speculative_tokens=num_speculative_tokens,
+    )
+
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,), (1,)),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+        num_speculative_blocks=num_speculative_blocks,
+    )
+
+    builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+
+    expected_cols = (max_model_len // block_size) + num_speculative_blocks
+    assert builder.state_indices_tensor_d.shape == (max_num_seqs, expected_cols)
+
+
+def test_block_idx_cudagraph_capture_padded_by_num_reqs():
+    """Regression test for https://github.com/vllm-project/vllm/issues/39809
+    bug 2: with mamba_cache_mode='all' and spec decode, _update_metadata_for
+    _cudagraph_capture must slice block_idx_last_{scheduled,computed}_token
+    by the request count (padded_bs == num_reqs), not by num_decode_tokens.
+    Past num_decodes, the slice must be zero-filled."""
+
+    block_size = 16
+    max_model_len = 256
+    max_num_seqs = 8
+    num_speculative_tokens = 1
+    device = torch.device("cpu")
+
+    vllm_config = _make_vllm_config(
+        max_model_len,
+        max_num_seqs,
+        num_speculative_tokens=num_speculative_tokens,
+    )
+
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,), (1,)),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+        num_speculative_blocks=2,
+    )
+
+    builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+
+    builder.block_idx_last_scheduled_token.fill_(-1)
+    builder.block_idx_last_computed_token.fill_(-1)
+
+    num_decodes = 2
+    num_reqs = 3
+    num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
+    seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
+    block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
+    state_indices_d = torch.zeros(
+        (num_decodes, builder.state_indices_tensor_d.shape[1]),
+        dtype=torch.int32,
+        device=device,
+    )
+    query_start_loc_d = torch.arange(
+        num_decodes + 1, dtype=torch.int32, device=device
+    ) * (1 + num_speculative_tokens)
+    num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
+
+    metadata = BaseMambaAttentionMetadata(
+        num_prefills=0,
+        num_prefill_tokens=0,
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decode_tokens,
+        num_reqs=num_reqs,
+        has_initial_states_p=None,
+        query_start_loc_p=None,
+        num_computed_tokens_p=None,
+        state_indices_tensor_p=None,
+        state_indices_tensor_d=state_indices_d,
+        query_start_loc_d=query_start_loc_d,
+        num_accepted_tokens=num_accepted_tokens,
+        block_idx_last_scheduled_token=block_idx_vals,
+        block_idx_first_scheduled_token_p=None,
+        block_idx_last_computed_token=block_idx_vals,
+        block_idx_last_scheduled_token_prev_step=None,
+        seq_lens=seq_lens,
+    )
+
+    out = builder._update_metadata_for_cudagraph_capture(metadata)
+
+    assert out.block_idx_last_scheduled_token.shape == (num_reqs,)
+    assert out.block_idx_last_computed_token.shape == (num_reqs,)
+    torch.testing.assert_close(
+        out.block_idx_last_scheduled_token[:num_decodes], block_idx_vals
+    )
+    torch.testing.assert_close(
+        out.block_idx_last_computed_token[:num_decodes], block_idx_vals
+    )
+    assert torch.all(out.block_idx_last_scheduled_token[num_decodes:] == 0)
+    assert torch.all(out.block_idx_last_computed_token[num_decodes:] == 0)
+
+
+def test_block_idx_prev_step_persistent_buffer_allocated():
+    """With mamba_cache_mode='all' + spec decode, the builder must allocate
+    block_idx_last_scheduled_token_prev_step as a persistent buffer with the
+    same shape as the existing block_idx_last_{scheduled,computed}_token
+    buffers, so cudagraph capture records a stable pointer for the prev-step
+    input anchor consumed by mamba_mixer2's input gather."""
+    block_size = 16
+    max_model_len = 256
+    max_num_seqs = 8
+    num_speculative_tokens = 1
+    device = torch.device("cpu")
+
+    vllm_config = _make_vllm_config(
+        max_model_len,
+        max_num_seqs,
+        num_speculative_tokens=num_speculative_tokens,
+    )
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,), (1,)),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+        num_speculative_blocks=2,
+    )
+    builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+
+    assert hasattr(builder, "block_idx_last_scheduled_token_prev_step")
+    assert builder.block_idx_last_scheduled_token_prev_step.shape == (max_num_seqs,)
+    assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32
+
+
+def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode():
+    """Without spec decode, the prev-step buffer is unused and must not be
+    allocated — the input anchor reduces to last_computed_token."""
+    block_size = 16
+    max_model_len = 256
+    max_num_seqs = 8
+    device = torch.device("cpu")
+
+    vllm_config = _make_vllm_config(
+        max_model_len, max_num_seqs, num_speculative_tokens=0
+    )
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,), (1,)),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+    )
+    builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+
+    assert not hasattr(builder, "block_idx_last_scheduled_token_prev_step")
+
+
+def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer():
+    """_update_metadata_for_cudagraph_capture must copy the prev-step anchor
+    into the builder's persistent buffer (so cudagraph replay reads from the
+    same underlying memory), pad past num_decodes with zero, and return a
+    slice of the persistent buffer in the metadata."""
+    block_size = 16
+    max_model_len = 256
+    max_num_seqs = 8
+    num_speculative_tokens = 1
+    device = torch.device("cpu")
+
+    vllm_config = _make_vllm_config(
+        max_model_len,
+        max_num_seqs,
+        num_speculative_tokens=num_speculative_tokens,
+    )
+    spec = MambaSpec(
+        block_size=block_size,
+        shapes=((1,), (1,)),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="all",
+        num_speculative_blocks=2,
+    )
+    builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
+    builder.block_idx_last_scheduled_token.fill_(-1)
+    builder.block_idx_last_computed_token.fill_(-1)
+    builder.block_idx_last_scheduled_token_prev_step.fill_(-1)
+
+    num_decodes = 2
+    num_reqs = 3
+    num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
+    seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
+    block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
+    prev_step_vals = torch.tensor([2, 4], dtype=torch.int32, device=device)
+    state_indices_d = torch.zeros(
+        (num_decodes, builder.state_indices_tensor_d.shape[1]),
+        dtype=torch.int32,
+        device=device,
+    )
+    query_start_loc_d = torch.arange(
+        num_decodes + 1, dtype=torch.int32, device=device
+    ) * (1 + num_speculative_tokens)
+    num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
+
+    metadata = BaseMambaAttentionMetadata(
+        num_prefills=0,
+        num_prefill_tokens=0,
+        num_decodes=num_decodes,
+        num_decode_tokens=num_decode_tokens,
+        num_reqs=num_reqs,
+        has_initial_states_p=None,
+        query_start_loc_p=None,
+        num_computed_tokens_p=None,
+        state_indices_tensor_p=None,
+        state_indices_tensor_d=state_indices_d,
+        query_start_loc_d=query_start_loc_d,
+        num_accepted_tokens=num_accepted_tokens,
+        block_idx_last_scheduled_token=block_idx_vals,
+        block_idx_first_scheduled_token_p=None,
+        block_idx_last_computed_token=block_idx_vals,
+        block_idx_last_scheduled_token_prev_step=prev_step_vals,
+        seq_lens=seq_lens,
+    )
+
+    out = builder._update_metadata_for_cudagraph_capture(metadata)
+
+    # Output field exists and is identity-shared with the persistent buffer.
+    assert out.block_idx_last_scheduled_token_prev_step is not None
+    assert (
+        out.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
+        == builder.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
+    ), (
+        "prev-step buffer must live in the builder's persistent buffer, not "
+        "in the caller-provided tensor"
+    )
+
+    # Padded by num_reqs (not num_decode_tokens) — same fix as bug 2 for the
+    # other block_idx_* fields.
+    assert out.block_idx_last_scheduled_token_prev_step.shape == (num_reqs,)
+
+    # First num_decodes values: input values copied through.
+    torch.testing.assert_close(
+        out.block_idx_last_scheduled_token_prev_step[:num_decodes],
+        prev_step_vals,
+    )
+
+    # Tail values past num_decodes: zero-filled padding for cudagraph capture.
+    assert torch.all(out.block_idx_last_scheduled_token_prev_step[num_decodes:] == 0)
@@ -24,8 +24,16 @@ from vllm.v1.attention.backend import (
    AttentionType,
    CommonAttentionMetadata,
 )
+from vllm.v1.attention.backends.mamba_attn import (
+    BaseMambaAttentionMetadata,
+    BaseMambaAttentionMetadataBuilder,
+)
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, FullAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    EncoderOnlyAttentionSpec,
+    FullAttentionSpec,
+    MambaSpec,
+)


@dataclass
@@ -376,3 +384,34 @@ full_cg_backend_configs = {
        },
    ),
 }
+
+
+class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]):
+    """Minimal concrete subclass for testing (base class is ABC)."""
+
+    metadata_cls = BaseMambaAttentionMetadata
+
+    @classmethod
+    def build_mamba_metadata(
+        cls,
+        vllm_config: VllmConfig,
+        seq_lens: list[int],
+        query_lens: list[int],
+        is_prefilling: list[bool],
+        *,
+        device: torch.device | None = None,
+    ) -> BaseMambaAttentionMetadata:
+        block_size = vllm_config.cache_config.block_size
+        device = device or torch.device("cpu")
+        mamba_spec = MambaSpec(
+            block_size=block_size, shapes=((1,), (1,)), dtypes=(torch.float32,)
+        )
+        builder = cls(mamba_spec, ["layer0"], vllm_config, device)
+        batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)
+        common_metadata = create_common_attn_metadata(
+            batch_spec, block_size=block_size, device=device, arange_block_indices=True
+        )
+        common_metadata = common_metadata.replace(
+            is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)
+        )
+        return builder.build(0, common_metadata)
@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for the breakable cudagraph primitives.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+
+import pytest
+import torch
+
+os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
+
+
+@pytest.fixture(autouse=True)
+def _reset_breakable_tls():
+    """Defensively clear thread-local capture state between tests so a
+    failure in one test can't leak "nested capture" errors into the next."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    BreakableCUDAGraphCapture._tls.active = None
+    yield
+    BreakableCUDAGraphCapture._tls.active = None
+
+
+@pytest.fixture
+def cuda_capture_stream():
+    """A non-default CUDA stream suitable for cudagraph capture.
+
+    ``CUDAGraph.capture_begin`` refuses to capture from the default
+    stream, so all capture-using tests need to run under
+    ``torch.cuda.stream(...)`` for a separate stream.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        yield stream
+    torch.cuda.current_stream().wait_stream(stream)
+
+
+# ---------------------------------------------------------------------------
+# eager_break_during_capture: outside capture
+# ---------------------------------------------------------------------------
+
+
+def test_decorator_passthrough_outside_capture():
+    from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+
+    calls = []
+
+    @eager_break_during_capture
+    def f(x):
+        calls.append(x)
+        return x * 2
+
+    assert f(3) == 6
+    assert calls == [3]
+
+
+# ---------------------------------------------------------------------------
+# BreakableCUDAGraphCapture: thread-local + nested rejection
+# ---------------------------------------------------------------------------
+
+
+def test_current_is_none_when_inactive():
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    assert BreakableCUDAGraphCapture.current() is None
+    assert BreakableCUDAGraphCapture.is_active() is False
+
+
+def test_thread_local_active_during_context(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        assert BreakableCUDAGraphCapture.current() is cap
+        assert BreakableCUDAGraphCapture.is_active() is True
+    assert BreakableCUDAGraphCapture.current() is None
+
+
+def test_nested_capture_raises(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    outer = BreakableCUDAGraphCapture()
+    inner = BreakableCUDAGraphCapture()
+    with outer, pytest.raises(RuntimeError, match="Nested.*not supported"), inner:
+        pass
+
+
+def test_active_state_isolated_across_threads(cuda_capture_stream):
+    """Verify the thread-local 'active capture' slot is per-thread.
+
+    We don't run concurrent captures here -- CUDA only supports one
+    in-flight capture per stream and we keep tests cheap. We just check
+    that the worker thread sees its own slot as None while the main
+    thread has a capture active.
+    """
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    worker_view: dict[str, BreakableCUDAGraphCapture | None] = {}
+
+    def worker():
+        worker_view["state"] = BreakableCUDAGraphCapture.current()
+
+    main_cap = BreakableCUDAGraphCapture()
+    with main_cap:
+        # Main thread has a live capture.
+        assert BreakableCUDAGraphCapture.current() is main_cap
+        t = threading.Thread(target=worker)
+        t.start()
+        t.join()
+
+    # Worker thread saw None -- thread-local separation works.
+    assert worker_view["state"] is None
+    # Main thread's slot is cleared on exit.
+    assert BreakableCUDAGraphCapture.current() is None
+
+
+# ---------------------------------------------------------------------------
+# Segment list construction
+# ---------------------------------------------------------------------------
+
+
+def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    x = torch.zeros(4, device="cuda")
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(1.0)
+    assert len(cap.segments) == 1
+    assert cap.num_graphs == 1
+    assert cap.num_eager_breaks == 0
+
+
+def test_add_eager_creates_alternating_graph_eager_graph(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    x = torch.zeros(4, device="cuda")
+    counter = {"eager_calls": 0}
+
+    def eager_step():
+        counter["eager_calls"] += 1
+        x.add_(10.0)
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(1.0)
+        cap.add_eager(eager_step)
+        x.add_(1.0)
+        cap.add_eager(eager_step)
+        x.add_(1.0)
+    # 3 graph segments + 2 eager segments, interleaved as G E G E G.
+    assert len(cap.segments) == 5
+    assert cap.num_graphs == 3
+    assert cap.num_eager_breaks == 2
+    # Eager fn is stored as-is in the segment list, so we can confirm
+    # the alternation pattern by identity check.
+    assert cap.segments[1] is eager_step
+    assert cap.segments[3] is eager_step
+    assert counter["eager_calls"] == 2  # only the in-capture invocation
+
+
+# ---------------------------------------------------------------------------
+# Capture vs eager numerical equivalence
+# ---------------------------------------------------------------------------
+
+
+def test_capture_replay_matches_eager_simple(cuda_capture_stream):
+    """Verify that replay reproduces the same end-state as a single eager
+    forward, with an eager break in the middle.
+
+    Note: during capture, the *captured* kernels are recorded but NOT
+    executed (that's CUDA-graph semantics). Only the eager segments
+    actually mutate state at capture time. So we check correctness after
+    ``replay()``, not after ``with cap:`` exits.
+    """
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    x = torch.zeros(8, device="cuda")
+    log: list[str] = []
+
+    def eager_break_op():
+        x.mul_(2.0)
+        log.append("eager")
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(1.0)  # recorded into graph[0]
+        cap.add_eager(eager_break_op)  # runs eagerly: x *= 2
+        x.add_(5.0)  # recorded into graph[1]
+
+    # Capture-time: graph kernels were recorded only; eager segment ran
+    # once on x == 0, leaving x == 0.
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.zeros(8, device="cuda"))
+    assert log == ["eager"]
+
+    # Replay with a fresh input: 10 -> 11 -> 22 -> 27.
+    x.fill_(10.0)
+    cap.replay()
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.full((8,), 27.0, device="cuda"))
+    assert log == ["eager", "eager"]
+
+    # Replay again with another input: 100 -> 101 -> 202 -> 207.
+    x.fill_(100.0)
+    cap.replay()
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.full((8,), 207.0, device="cuda"))
+    assert log == ["eager", "eager", "eager"]
+
+
+def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream):
+    """Verify @eager_break_during_capture correctly routes through
+    add_eager when inside a capture context, and runs straight through
+    when there's no active capture."""
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+
+    @eager_break_during_capture
+    def attention_like(t: torch.Tensor) -> None:
+        # In-place double; stands in for "real" attention work.
+        t.mul_(2.0)
+
+    x = torch.zeros(4, device="cuda")
+
+    # Outside capture: decorator should just call through.
+    x.fill_(3.0)
+    attention_like(x)
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.full((4,), 6.0, device="cuda"))
+
+    # Inside capture: decorator should split the graph. Only the eager
+    # segment actually mutates state during capture.
+    x.fill_(0.0)
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(5.0)  # recorded
+        attention_like(x)  # eager: x *= 2 (on x == 0, no-op)
+        x.add_(1.0)  # recorded
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.zeros(4, device="cuda"))
+    # 2 graph segments + 1 eager segment, ordered G E G; the arithmetic
+    # equivalence check below verifies the ordering.
+    assert len(cap.segments) == 3
+    assert cap.num_graphs == 2
+    assert cap.num_eager_breaks == 1
+
+    # Replay: 2 -> 7 -> 14 -> 15.
+    x.fill_(2.0)
+    cap.replay()
+    torch.accelerator.synchronize()
+    assert torch.equal(x, torch.full((4,), 15.0, device="cuda"))
+
+
+# ---------------------------------------------------------------------------
+# Replay ordering
+# ---------------------------------------------------------------------------
+
+
+def test_replay_invokes_eager_segments_in_order(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    log: list[str] = []
+    x = torch.zeros(1, device="cuda")
+
+    def make_eager(name):
+        def step():
+            log.append(name)
+            x.add_(1.0)
+
+        return step
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(1.0)
+        cap.add_eager(make_eager("A"))
+        x.add_(1.0)
+        cap.add_eager(make_eager("B"))
+        x.add_(1.0)
+        cap.add_eager(make_eager("C"))
+        x.add_(1.0)
+
+    # Capture-time invocation order
+    assert log == ["A", "B", "C"]
+
+    log.clear()
+    cap.replay()
+    torch.accelerator.synchronize()
+    assert log == ["A", "B", "C"]
+
+
+# ---------------------------------------------------------------------------
+# Capture cleanup releases thread-local even if body raises
+# ---------------------------------------------------------------------------
+
+
+def test_exception_in_body_clears_active(cuda_capture_stream):
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    cap = BreakableCUDAGraphCapture()
+    with pytest.raises(RuntimeError, match="boom"), cap:
+        raise RuntimeError("boom")
+
+    # active must be reset even after an exception inside the body
+    assert BreakableCUDAGraphCapture.current() is None
+
+
+# ---------------------------------------------------------------------------
+# Nested decorated ops: inner op must not trigger a recursive eager break
+# ---------------------------------------------------------------------------
+
+
+def test_nested_decorated_op_runs_inline(cuda_capture_stream):
+    """A decorated op invoked from inside another decorated op's eager
+    body must execute inline -- starting a second eager break mid-flight
+    corrupts the segment state and explodes ``_begin_segment``'s assert.
+
+    This mirrors the deepseek_v4_attention case where the outer attention
+    op's impl internally dispatches sparse_attn_indexer (also decorated).
+    """
+    from vllm.compilation.breakable_cudagraph import (
+        BreakableCUDAGraphCapture,
+        eager_break_during_capture,
+    )
+
+    x = torch.zeros(4, device="cuda")
+    inner_calls = 0
+
+    @eager_break_during_capture
+    def inner_op(t: torch.Tensor) -> None:
+        nonlocal inner_calls
+        inner_calls += 1
+        t.add_(1.0)
+
+    @eager_break_during_capture
+    def outer_op(t: torch.Tensor) -> None:
+        # outer body calls another decorated op -- this is the case that
+        # used to assert in _begin_segment.
+        inner_op(t)
+        t.add_(10.0)
+
+    cap = BreakableCUDAGraphCapture()
+    with cap:
+        x.add_(2.0)  # recorded in graph[0]
+        outer_op(x)  # one eager break, inner runs inline
+        x.add_(100.0)  # recorded in graph[1]
+
+    # Exactly one eager break (the outer); inner must NOT add a second.
+    assert cap.num_graphs == 2
+    assert cap.num_eager_breaks == 1
+    assert inner_calls == 1  # only the capture-time invocation
+
+    x.fill_(0.0)
+    cap.replay()
+    torch.accelerator.synchronize()
+    # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113
+    assert torch.equal(x, torch.full((4,), 113.0, device="cuda"))
+    assert inner_calls == 2  # replay invokes the outer's lambda again
@@ -364,26 +364,34 @@ def get_fake_process_mamba_fn(
    def fake_post_process_mamba_fn(
        scheduler_output: SchedulerOutput,
        kv_cache_config: KVCacheConfig,
+        cache_config: CacheConfig,
        input_batch: GPUInputBatch,
        requests: dict[str, CachedRequestState],
        mamba_state_idx: dict[str, int],
-        forward_context: dict[str, Any],
-        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
-        copy_bufs: mamba_utils.MambaCopyBuffers,
+        num_spec_tokens: int,
+        num_reqs: int,
+        *,
+        forward_context: dict[str, Any] | None = None,
+        mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...] | None = None,
+        copy_bufs: mamba_utils.MambaCopyBuffers | None = None,
    ):
        nonlocal copy_info
        copy_info = None
        ret = original_post_process_mamba_fn(
            scheduler_output,
            kv_cache_config,
+            cache_config,
            input_batch,
            requests,
            mamba_state_idx,
-            forward_context,
-            mamba_state_copy_funcs,
-            copy_bufs,
+            num_spec_tokens,
+            num_reqs,
+            forward_context=forward_context,
+            mamba_state_copy_funcs=mamba_state_copy_funcs,
+            copy_bufs=copy_bufs,
        )
        if cur_step_action is not None:
+            assert forward_context is not None
            check_copy_info(
                cur_step_action.postprocess_copy_idx,
                kv_cache_config,
@@ -3,6 +3,10 @@ set -euo pipefail

 # Utility to run integration tests sequentially with varying TP configurations.
 SCRIPT="v1/kv_connector/nixl_integration/run_accuracy_test.sh"
+IMPORT_CANARY="v1/kv_connector/nixl_integration/test_nixl_imports.py"
+
+echo "=== Running NIXL import canary ==="
+python3 -m pytest -s -x "${IMPORT_CANARY}"

 # Define test configurations
 tp_configs=(
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NIXL import canaries for CUDA wheel selection."""
+
+import importlib
+import importlib.metadata as metadata
+import pathlib
+import subprocess
+import sys
+
+import pytest
+import torch
+
+
+def _print_distribution_version(package_name: str) -> None:
+    try:
+        version = metadata.version(package_name)
+    except metadata.PackageNotFoundError:
+        version = "not installed"
+    print(f"{package_name}: {version}")
+
+
+@pytest.mark.skipif(torch.version.cuda is None, reason="CUDA NIXL EP canary")
+def test_nixl_and_nixl_ep_imports() -> None:
+    """Verify both core NIXL and the NIXL EP extension import successfully."""
+    print(f"torch cuda: {torch.version.cuda}")
+    for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
+        _print_distribution_version(package_name)
+
+    nixl = importlib.import_module("nixl")
+    print(f"nixl: {nixl.__file__}")
+
+    # Exercise the core NIXL bindings used by NixlConnector.
+    importlib.import_module("nixl._api")
+    importlib.import_module("nixl._bindings")
+
+    # Exercise the NIXL EP extension used by fused MoE expert parallelism.
+    nixl_ep = importlib.import_module("nixl_ep")
+    print(f"nixl_ep: {nixl_ep.__file__}")
+
+    assert nixl_ep.__file__ is not None
+    extension_dir = pathlib.Path(nixl_ep.__file__).parent
+    extension_files = sorted(extension_dir.glob("nixl_ep_cpp*.so"))
+    assert extension_files, f"No nixl_ep_cpp extension found in {extension_dir}"
+
+    extension_file = extension_files[0]
+    completed = subprocess.run(
+        ["ldd", str(extension_file)],
+        capture_output=True,
+        check=False,
+        text=True,
+    )
+    print(completed.stdout)
+    if completed.stderr:
+        print(completed.stderr, file=sys.stderr)
+
+    assert completed.returncode == 0
+    if torch.version.cuda is not None:
+        cuda_major = torch.version.cuda.split(".", maxsplit=1)[0]
+        expected_cudart = f"libcudart.so.{cuda_major}"
+        assert expected_cudart in completed.stdout
+        assert f"{expected_cudart} => not found" not in completed.stdout
@@ -258,9 +258,12 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
+    # With sync scheduling, all-finished flush fires within this run.
+    # With async scheduling, the finish is delayed so flush fires later.
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        expected_stored=(0, 1, 2),
+        expected_flushed=(0, 1, 2) if not async_scheduling else (),
    )

    # start a request to load the first block, but don't complete
@@ -325,6 +328,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        expected_stored=(0, 1, 2),
+        expected_flushed=(0, 1, 2) if not async_scheduling else (),
    )

    # start a request to load the first block, but don't complete
@@ -351,6 +355,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
    runner.run(
        decoded_tokens=[],
        expected_loaded=(0, 1, 2),
+        expected_flushed=(0, 1, 2),
    )

    # assert request is deleted
@@ -766,7 +771,11 @@ def test_do_remote_decode_stores_all_blocks(request_runner, async_scheduling: bo
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=(0, 1, 2),
+        expected_flushed=(0, 1, 2) if not async_scheduling else (),
+    )

    # Reset GPU prefix cache so the next request must load from CPU.
    runner.scheduler.reset_prefix_cache()
@@ -831,8 +840,13 @@ def test_fence_at_update_state_after_alloc(request_runner):
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
-    assert runner.connector_scheduler._block_id_to_pending_jobs
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        complete_transfers=False,
+        expected_stored=(0,),
+        expected_flushed=(0,),
+    )
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}

    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * 4)
@@ -843,8 +857,6 @@ def test_fence_at_update_state_after_alloc(request_runner):
    runner.run(
        decoded_tokens=[],
        complete_transfers=False,
-        expected_stored=(0,),
-        expected_flushed=(0,),
    )
    assert runner.connector_scheduler._block_id_to_pending_jobs == {}

@@ -864,8 +876,13 @@ def test_fence_at_build_store_jobs(request_runner):
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID], complete_transfers=False)
-    assert runner.connector_scheduler._block_id_to_pending_jobs
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        complete_transfers=False,
+        expected_stored=(0,),
+        expected_flushed=(0,),
+    )
+    assert runner.connector_scheduler._block_id_to_pending_jobs == {}

    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[1] * 4)
@@ -875,8 +892,6 @@ def test_fence_at_build_store_jobs(request_runner):
    )
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored=(0,),
-        expected_flushed=(0,),
    )
    assert runner.connector_scheduler._block_id_to_pending_jobs == {}

@@ -921,6 +936,32 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
    assert runner.manager.complete_store.call_count == 0


+def test_flush_all_jobs_when_no_requests_remain(request_runner):
+    """When all tracked requests are finished, build_connector_meta flushes
+    all pending jobs since there will be no future step to complete them."""
+    block_size = 4
+    block_size_factor = 1
+    offloaded_block_size = block_size * block_size_factor
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=False,
+        block_size_factor=block_size_factor,
+    )
+
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        complete_transfers=False,
+        expected_stored=(0,),
+        expected_flushed=(0,),
+    )
+
+
@pytest.mark.parametrize("async_scheduling", [True, False])
 def test_reset_cache(request_runner, async_scheduling: bool):
    """reset_cache flushes in-flight loads, calls manager.reset_cache(), resets
@@ -942,7 +983,11 @@ def test_reset_cache(request_runner, async_scheduling: bool):
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2))
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=(0, 1, 2),
+        expected_flushed=(0, 1, 2) if not async_scheduling else (),
+    )

    # Reset GPU prefix cache then start a request that loads from CPU.
    # Leave the load in-flight so that reset_cache must flush it.
@@ -348,10 +348,14 @@ class RequestRunner:
    def _parse_transfers(self):
        for transfer_spec in self.offloading_spec.get_flushed_transfers():
            src_spec, dst_spec = transfer_spec
-            assert isinstance(src_spec, GPULoadStoreSpec)
-
-            for block_id in src_spec.block_ids:
-                self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
+            if isinstance(src_spec, GPULoadStoreSpec):
+                # store flush
+                for block_id in src_spec.block_ids:
+                    self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])
+            else:
+                # load flush
+                for block_id in dst_spec.block_ids:
+                    self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()])

        block_size_factor = self.block_size_factor

@@ -1,23 +1,36 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import json
+import logging
+import math
+import sys
 import threading
+import types
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

+import pytest
 import torch

+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
+    rdma_utils,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
    worker,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
    ChunkedTokenDatabase,
    KeyMetadata,
+    LoadSpec,
    ReqMeta,
 )


 def _make_store_sending_thread(
    store: MagicMock,
+    *,
+    replicate_config: object | None = None,
 ) -> worker.KVCacheStoreSendingThread:
    token_database = ChunkedTokenDatabase(
        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
@@ -32,11 +45,55 @@ def _make_store_sending_thread(
        put_step=1,
        kv_role="kv_producer",
        ready_event=threading.Event(),
+        replicate_config=replicate_config,
    )
    thread.request_queue.task_done = MagicMock()
    return thread


+def _make_store_recving_thread(
+    store: MagicMock,
+    *,
+    disk_offload_buffer_budget_bytes: int | None = None,
+) -> worker.KVCacheStoreRecvingThread:
+    token_database = ChunkedTokenDatabase(
+        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
+    )
+    token_database.set_kv_caches_base_addr([0x1000])
+    token_database.set_block_len([256])
+    thread = worker.KVCacheStoreRecvingThread(
+        store=store,
+        token_database=token_database,
+        block_size=16,
+        tp_rank=0,
+        ready_event=threading.Event(),
+        disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
+    )
+    thread.request_queue.task_done = MagicMock()
+    return thread
+
+
+def _make_load_req(
+    req_id: str,
+    block_hashes: list[bytes],
+    *,
+    token_len: int,
+    vllm_cached_tokens: int = 0,
+) -> ReqMeta:
+    return ReqMeta(
+        req_id=req_id,
+        token_len_chunk=token_len,
+        block_ids=list(range(len(block_hashes))),
+        block_hashes=block_hashes,
+        load_spec=LoadSpec(
+            vllm_cached_tokens=vllm_cached_tokens,
+            kvpool_cached_tokens=token_len,
+            can_load=True,
+            token_len=token_len,
+        ),
+    )
+
+
 def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
    return ReqMeta(
        req_id=req_id,
@@ -48,6 +105,233 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
    )


+_DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
+_DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
+_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
+_DISK_OFFLOAD_BUDGET_FOR_SPLIT = math.ceil(
+    2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES / _DISK_OFFLOAD_USABLE_BUDGET_RATIO
+)  # Allows two 256-byte chunks but not the third.
+_DISK_OFFLOAD_BUDGET_TOO_SMALL = (
+    _DISK_OFFLOAD_SINGLE_KEY_BYTES - 1
+)  # Smaller than a single 256-byte chunk.
+
+
+class _FakeKVTransferConfig:
+    def __init__(
+        self,
+        *,
+        kv_role: str = "kv_both",
+        extra_config: dict[str, object] | None = None,
+    ) -> None:
+        self.kv_role = kv_role
+        self.kv_connector_extra_config = extra_config or {}
+
+    def get_from_extra_config(self, key: str, default: object) -> object:
+        return self.kv_connector_extra_config.get(key, default)
+
+
+class _FakeModelConfig:
+    model = "test-model"
+    use_mla = False
+
+    def get_num_layers(self, parallel_config) -> int:
+        return 1
+
+    def get_total_num_kv_heads(self) -> int:
+        return 1
+
+
+def _make_vllm_config(
+    *, extra_config: dict[str, object] | None = None
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        model_config=_FakeModelConfig(),
+        parallel_config=SimpleNamespace(
+            pipeline_parallel_size=1,
+            rank=0,
+        ),
+        kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
+        cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),
+        kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
+    )
+
+
+def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str:
+    config_path = tmp_path / "mooncake_config.json"
+    config_path.write_text(json.dumps(config), encoding="utf-8")
+    return str(config_path)
+
+
+def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
+    class FakeReplicateConfig:
+        def __init__(self) -> None:
+            self.preferred_segment = ""
+
+    fake_store_module = types.ModuleType("mooncake.store")
+    fake_store_module.MooncakeDistributedStore = lambda: store_instance  # type: ignore[attr-defined]
+    fake_store_module.ReplicateConfig = FakeReplicateConfig  # type: ignore[attr-defined]
+    fake_mooncake_module = types.ModuleType("mooncake")
+    fake_mooncake_module.store = fake_store_module  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
+    monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module)
+    return FakeReplicateConfig
+
+
+def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
+    single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
+    monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
+    monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
+    monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
+    monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
+
+
+def test_default_local_buffer_size_matches_pr40900():
+    """PR-40900 shipped a 4 GiB default for local_buffer_size; the dual-mode
+    patch preserves it (and the JSON key) so unchanged PR-40900 configs work."""
+    assert worker.DEFAULT_LOCAL_BUFFER_SIZE == 4 * 1024**3
+
+
+def test_get_requester_local_hostname_prefers_override(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
+
+    assert rdma_utils.get_requester_local_hostname("10.0.0.7") == "worker-a:50053"
+
+
+def test_get_configured_preferred_segment_returns_explicit_override():
+    assert (
+        rdma_utils.get_configured_preferred_segment(
+            {"preferred_segment": "10.0.0.7:50053"}
+        )
+        == "10.0.0.7:50053"
+    )
+
+
+def test_get_configured_preferred_segment_prefers_explicit_over_env(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
+
+    assert (
+        rdma_utils.get_configured_preferred_segment(
+            {"preferred_segment": "10.0.0.7:50053"}
+        )
+        == "10.0.0.7:50053"
+    )
+
+
+def test_get_configured_preferred_segment_returns_env_override(monkeypatch):
+    monkeypatch.setenv("MOONCAKE_PREFERRED_SEGMENT", "10.0.0.8:50053")
+
+    assert rdma_utils.get_configured_preferred_segment({}) == "10.0.0.8:50053"
+
+
+def test_get_configured_preferred_segment_rejects_empty_override():
+    with pytest.raises(ValueError, match="preferred_segment"):
+        rdma_utils.get_configured_preferred_segment({"preferred_segment": "  "})
+
+
+def test_get_configured_worker_rnic_prefers_explicit_device_name(monkeypatch):
+    store_config = worker.MooncakeStoreConfig(
+        metadata_server="",
+        local_buffer_size=1,
+        protocol="rdma",
+        device_name="rocep139s0",
+        master_server_address="",
+    )
+
+    assert (
+        rdma_utils.get_configured_worker_rnic(
+            protocol=store_config.protocol,
+            configured_device=store_config.device_name,
+        )
+        == "rocep139s0"
+    )
+
+
+def test_get_configured_worker_rnic_selects_device_from_explicit_csv(monkeypatch):
+    monkeypatch.setattr(
+        rdma_utils,
+        "get_current_physical_gpu_index",
+        lambda: 1,
+    )
+    store_config = worker.MooncakeStoreConfig(
+        metadata_server="",
+        local_buffer_size=1,
+        protocol="rdma",
+        device_name="rocep139s0,rocep140s0",
+        master_server_address="",
+    )
+
+    assert (
+        rdma_utils.get_configured_worker_rnic(
+            protocol=store_config.protocol,
+            configured_device=store_config.device_name,
+        )
+        == "rocep140s0"
+    )
+
+
+def test_get_configured_worker_rnic_warns_and_returns_empty_for_rdma_with_no_device(
+    caplog, monkeypatch
+):
+    """No device configured + protocol=rdma → emit a clear warning and return ""
+    so the C++ side handles auto-selection. There is no Python-side fallback."""
+    monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
+    with caplog.at_level(logging.WARNING):
+        result = rdma_utils.get_configured_worker_rnic(
+            protocol="rdma",
+            configured_device="",
+        )
+    assert result == ""
+    warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+    assert any("No RDMA devices specified" in r.message for r in warnings), (
+        f"expected fallback warning, got {[r.message for r in warnings]}"
+    )
+
+
+def test_get_configured_worker_rnic_silent_for_tcp_with_no_device(caplog, monkeypatch):
+    """protocol=tcp + no device → return "" silently (no RDMA, no warning)."""
+    monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
+    with caplog.at_level(logging.WARNING):
+        result = rdma_utils.get_configured_worker_rnic(
+            protocol="tcp",
+            configured_device="",
+        )
+    assert result == ""
+    warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+    assert not any("RDMA" in r.message for r in warnings), (
+        "did not expect RDMA warning for tcp protocol, got "
+        f"{[r.message for r in warnings]}"
+    )
+
+
+def test_get_configured_worker_rnic_rejects_short_explicit_csv(monkeypatch):
+    monkeypatch.setattr(
+        rdma_utils,
+        "get_current_physical_gpu_index",
+        lambda: 2,
+    )
+    with pytest.raises(ValueError, match="does not cover local GPU 2"):
+        rdma_utils.get_configured_worker_rnic(
+            protocol="rdma",
+            configured_device="rocep139s0,rocep140s0",
+        )
+
+
+class _ReplicaDesc:
+    def __init__(self, tier: str):
+        self.tier = tier
+
+    def is_memory_replica(self) -> bool:
+        return self.tier == "memory"
+
+    def is_disk_replica(self) -> bool:
+        return self.tier == "disk"
+
+    def is_local_disk_replica(self) -> bool:
+        return self.tier == "disk"
+
+
 def test_store_sending_thread_skips_request_during_cpu_pressure():
    store = MagicMock()
    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -105,6 +389,372 @@ def test_store_sending_thread_only_skips_on_no_available_handle():
    assert store.batch_put_from_multi_buffers.call_count == 2


+def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set():
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.return_value = [256, 256]
+    replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053")
+    thread = _make_store_sending_thread(store, replicate_config=replicate_config)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 1
+    call_args = store.batch_put_from_multi_buffers.call_args.args
+    assert len(call_args) == 4
+    assert call_args[3] is replicate_config
+
+
+def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_segment():  # noqa: E501
+    """Without a preferred_segment the SendingThread still forwards a
+    (default-constructed) ReplicateConfig so the C++ side always sees a
+    well-defined config object."""
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.return_value = [256, 256]
+    replicate_config = SimpleNamespace()
+    thread = _make_store_sending_thread(store, replicate_config=replicate_config)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 1
+    call_args = store.batch_put_from_multi_buffers.call_args.args
+    assert len(call_args) == 4
+    assert call_args[3] is replicate_config
+
+
+def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes():
+    assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288
+
+
+def test_recv_thread_uses_single_batch_when_no_disk_offload_budget(monkeypatch):
+    monkeypatch.delenv("VLLM_MOONCAKE_STORE_TIER_LOG", raising=False)
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
+    thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+    keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args
+    assert keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    assert sizes == [[256], [256], [256]]
+    store.batch_get_replica_desc.assert_not_called()
+
+
+def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):
+    monkeypatch.setenv("VLLM_MOONCAKE_STORE_TIER_LOG", "1")
+    caplog_vllm.set_level(logging.INFO, logger=worker.logger.name)
+
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, -10]
+    thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+    expected_keys = [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    store.batch_get_replica_desc.return_value = {
+        expected_keys[0]: [_ReplicaDesc("memory")],
+        expected_keys[1]: [_ReplicaDesc("disk")],
+        expected_keys[2]: [],
+    }
+
+    thread._handle_request(req)
+
+    assert store.batch_get_replica_desc.call_args.args == (expected_keys,)
+    assert store.method_calls[0][0] == "batch_get_replica_desc"
+    assert store.method_calls[1][0] == "batch_get_into_multi_buffers"
+
+    messages = [record.getMessage() for record in caplog_vllm.records]
+    assert any(
+        "Mooncake load tier summary" in message
+        and "req_id=req-a" in message
+        and "batch_keys=3" in message
+        and "memory_keys=1" in message
+        and "disk_keys=1" in message
+        and "unknown_keys=1" in message
+        and "success_keys=2" in message
+        and "failed_keys=1" in message
+        and "bytes_by_tier={'memory': 256, 'disk': 256, 'unknown': 0}" in message
+        for message in messages
+    )
+
+
+def test_recv_thread_uses_ratio_scaled_budget_for_first_pass_split():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.side_effect = [
+        [256],
+        [256],
+    ]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=2 * _DISK_OFFLOAD_SINGLE_KEY_BYTES,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1"],
+        token_len=32,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 2
+    first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
+    second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
+    assert first_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+    ]
+    assert second_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+    ]
+
+
+def test_recv_thread_splits_disk_offload_loads_by_budget():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.side_effect = [
+        [256, 256],
+        [256],
+    ]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 2
+
+    first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
+    second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
+    first_addrs = store.batch_get_into_multi_buffers.call_args_list[0].args[1]
+    second_addrs = store.batch_get_into_multi_buffers.call_args_list[1].args[1]
+    first_sizes = store.batch_get_into_multi_buffers.call_args_list[0].args[2]
+    second_sizes = store.batch_get_into_multi_buffers.call_args_list[1].args[2]
+    assert first_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+    ]
+    assert second_keys == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+    base_addr = thread.token_database.kv_caches_base_addr[0]
+    block_len = thread.token_database.block_len[0]
+    assert first_addrs == [[base_addr], [base_addr + block_len]]
+    assert second_addrs == [[base_addr + 2 * block_len]]
+    expected_size = block_len
+    assert first_sizes == [[expected_size], [expected_size]]
+    assert second_sizes == [[expected_size]]
+
+
+def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch():
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [-10, -10]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_SPLIT,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+
+
+def test_recv_thread_skips_split_when_budget_holds_all_keys():
+    """PR-36 removed the count-based split trigger; with budget for 3 keys,
+    all three should be requested in a single call."""
+    store = MagicMock()
+    store.batch_get_into_multi_buffers.return_value = [256, 256, 256]
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0", b"a1", b"a2"],
+        token_len=48,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 1
+    assert store.batch_get_into_multi_buffers.call_args_list[0].args[0] == [
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
+        "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
+    ]
+
+
+def test_recv_thread_reports_unsplittable_key_larger_than_budget():
+    store = MagicMock()
+    thread = _make_store_recving_thread(
+        store,
+        disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,
+    )
+
+    req = _make_load_req(
+        "req-a",
+        [b"a0"],
+        token_len=16,
+    )
+
+    thread._handle_request(req)
+
+    assert store.batch_get_into_multi_buffers.call_count == 0
+
+
+def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": "4gb",
+                "local_buffer_size": "64mb",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": True,
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert not hasattr(w, "_isolate_offload_resources")
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        4 * 1024 * 1024 * 1024,  # global_segment_size: "4gb" honored
+        64 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+
+
+def test_requester_worker_init_prefers_local_hostname_override(
+    tmp_path,
+    monkeypatch,
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv("MOONCAKE_REQUESTER_LOCAL_HOSTNAME", "worker-a:50053")
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "local_buffer_size": "64mb",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+    worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert store.setup.call_args.args[0] == "worker-a:50053"
+
+
+def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
+    tmp_path,
+    monkeypatch,
+):
+    """enable_offload=False zeroes out the disk budget so we don't generate
+    redundant owner GET-RPCs."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": False,
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    assert w.disk_offload_buffer_budget_bytes is None
+
+
+def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
+    tmp_path,
+    monkeypatch,
+):
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "protocol": "tcp",
+                "device_name": "",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+    w = worker.MooncakeStoreWorker(
+        _make_vllm_config(
+            extra_config={
+                "preferred_segment": "10.0.0.7:50053",
+            }
+        )
+    )
+
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
+
+
 # ---------------------------------------------------------------------------
 # Helpers for register_kv_caches tests
 # ---------------------------------------------------------------------------
@@ -147,8 +797,10 @@ def _make_bare_worker(
    w.tp_rank = 0
    w.put_step = 1
    w.enable_kv_events = False
+    w.disk_offload_buffer_budget_bytes = None
    w.kv_send_thread = None
    w.kv_recv_thread = None
+    w.store_replicate_config = SimpleNamespace()
    return w


@@ -298,3 +950,186 @@ def test_register_kv_caches_cross_layer_single_segment():

    assert w2.kv_caches_base_addr == w.kv_caches_base_addr
    assert w2.block_len == w.block_len
+
+
+# ---------------------------------------------------------------------------
+# Dual-mode (embedded / standalone-store) config validation tests
+# ---------------------------------------------------------------------------
+
+
+def _make_config(**overrides):
+    """Build a MooncakeStoreConfig with sensible defaults for validation tests.
+
+    Required dataclass fields are populated; callers override only the field
+    under test.
+    """
+    base = dict(
+        metadata_server="http://metadata/endpoint",
+        master_server_address="10.0.0.7:50051",
+        protocol="rdma",
+        device_name="mlx5_0",
+    )
+    base.update(overrides)
+    return worker.MooncakeStoreConfig(**base)
+
+
+def test_config_defaults_to_embedded():
+    """A JSON without explicit mode parses as embedded with 4 GiB segment."""
+    cfg = _make_config()
+    assert cfg.mode == "embedded"
+    assert cfg.global_segment_size == worker.DEFAULT_GLOBAL_SEGMENT_SIZE
+    assert cfg.local_buffer_size == worker.DEFAULT_LOCAL_BUFFER_SIZE
+    assert cfg.enable_offload is False
+
+
+def test_config_pr40900_unchanged(tmp_path):
+    """A literal PR-40900 config (no mode, no enable_offload, no preferred_segment)
+    parses without raising and resolves to embedded mode."""
+    config_path = _write_mooncake_config(
+        tmp_path,
+        {
+            "metadata_server": "http://metadata/endpoint",
+            "global_segment_size": "4GB",
+            "local_buffer_size": "4GB",
+            "protocol": "rdma",
+            "device_name": "mlx5_0",
+            "master_server_address": "10.0.0.7:50051",
+        },
+    )
+    cfg = worker.MooncakeStoreConfig.from_file(config_path)
+    assert cfg.mode == "embedded"
+    assert cfg.global_segment_size == 4 * 1024**3
+    assert cfg.local_buffer_size == 4 * 1024**3
+    assert cfg.enable_offload is False
+
+
+def test_config_embedded_rejects_zero_segment():
+    with pytest.raises(
+        ValueError, match=r"embedded mode requires global_segment_size > 0"
+    ):
+        _make_config(mode="embedded", global_segment_size=0)
+
+
+def test_config_standalone_store_rejects_nonzero_segment():
+    with pytest.raises(
+        ValueError,
+        match=r"standalone-store mode requires global_segment_size == 0",
+    ):
+        _make_config(mode="standalone-store", global_segment_size=4 * 1024**3)
+
+
+def test_config_standalone_store_accepts_zero_segment():
+    cfg = _make_config(mode="standalone-store", global_segment_size=0)
+    assert cfg.mode == "standalone-store"
+    assert cfg.global_segment_size == 0
+
+
+def test_config_unknown_mode():
+    with pytest.raises(ValueError, match=r"unknown Mooncake mode"):
+        _make_config(mode="something-else")
+
+
+def test_config_zero_local_buffer():
+    with pytest.raises(ValueError, match=r"local_buffer_size must be > 0"):
+        _make_config(local_buffer_size=0)
+
+
+# ---------------------------------------------------------------------------
+# End-to-end topology tests
+# Covers the two supported recipes:
+#   (A) standalone-store mode + disk offload (mode="standalone-store",
+#       segment=0, enable_offload=true, preferred_segment set)
+#   (B) embedded mode + CPU only      (mode default, segment>0,
+#       enable_offload=false, no preferred_segment)
+# ---------------------------------------------------------------------------
+
+
+def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):
+    """standalone-store + disk: global_segment_size=0, enable_offload=True,
+    preferred_segment set. Assert setup() positional args, ReplicateConfig
+    wiring, and that the disk-offload buffer budget is allocated."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "mode": "standalone-store",
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": 0,
+                "local_buffer_size": "1GB",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+                "enable_offload": True,
+            },
+        ),
+    )
+
+    w = worker.MooncakeStoreWorker(
+        _make_vllm_config(extra_config={"preferred_segment": "10.0.0.7:50053"})
+    )
+
+    # setup() receives global_segment_size=0 and the configured local buffer.
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        0,
+        1024 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+    # ReplicateConfig is built and carries the preferred_segment.
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
+    # Disk-offload staging budget is allocated (enable_offload=True).
+    assert w.disk_offload_buffer_budget_bytes is not None
+    assert w.disk_offload_buffer_budget_bytes > 0
+
+
+def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
+    """embedded + CPU-only: no mode key (defaults to embedded),
+    global_segment_size>0, enable_offload absent, no preferred_segment.
+    This is the PR-40900 baseline recipe."""
+    store = MagicMock()
+    store.setup.return_value = 0
+    fake_replicate_config_cls = _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setenv(
+        "MOONCAKE_CONFIG_PATH",
+        _write_mooncake_config(
+            tmp_path,
+            {
+                "metadata_server": "http://metadata/endpoint",
+                "global_segment_size": "4GB",
+                "local_buffer_size": "4GB",
+                "protocol": "rdma",
+                "device_name": "mlx5_0",
+                "master_server_address": "10.0.0.7:50051",
+            },
+        ),
+    )
+
+    w = worker.MooncakeStoreWorker(_make_vllm_config())
+
+    # setup() receives global_segment_size=4 GiB (rank contributes a segment).
+    assert store.setup.call_args.args == (
+        "10.0.0.7",
+        "http://metadata/endpoint",
+        4 * 1024 * 1024 * 1024,
+        4 * 1024 * 1024 * 1024,
+        "rdma",
+        "mlx5_0",
+        "10.0.0.7:50051",
+    )
+    # No preferred_segment — ReplicateConfig is default-constructed (so the
+    # preferred_segment field keeps its default value).
+    assert w.preferred_segment is None
+    assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
+    assert w.store_replicate_config.preferred_segment == ""
+    # No disk budget — enable_offload was absent (defaults to False).
+    assert w.disk_offload_buffer_budget_bytes is None
@@ -8,6 +8,7 @@ from unittest.mock import patch
 import pytest
 import torch

+from tests.v1.attention.utils import MockMambaBuilder
 from vllm import LLM, SamplingParams
 from vllm.config import KVTransferConfig
 from vllm.v1.core.single_type_kv_cache_manager import (
@@ -636,6 +637,30 @@ def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count):
    assert is_async is True


+@pytest.mark.cpu_test
+def test_mamba_n1_d_side_builds_decode_metadata():
+    req = create_request(num_tokens=10, do_remote_prefill=True)
+    sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
+
+    num_computed_tokens, is_async = sched.get_num_new_matched_tokens(
+        req, num_computed_tokens=0
+    )
+
+    assert num_computed_tokens == req.num_prompt_tokens - 1
+    assert is_async is True
+
+    vllm_config = create_vllm_config()
+    metadata = MockMambaBuilder.build_mamba_metadata(
+        vllm_config,
+        seq_lens=[req.num_prompt_tokens],
+        query_lens=[1],
+        is_prefilling=[True],
+    )
+
+    assert metadata.num_decodes == 1
+    assert metadata.num_prefills == 0
+
+
@pytest.mark.cpu_test
 def test_mamba_n1_p_side_truncation():
    """P-side: Mamba truncates prompt to N-1, sets max_tokens=1.
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Unit tests for TieringOffloadingManager and ExampleSecondaryTier.
+Unit tests for TieringOffloadingManager and ExampleSecondaryTierManager.

 These tests verify:
 1. Basic tiered offloading operations (store, load, lookup)
@@ -14,7 +14,6 @@ These tests verify:
 from collections.abc import Iterable
 from unittest.mock import MagicMock

-import numpy as np
 import pytest
 import torch

@@ -23,8 +22,7 @@ from vllm.v1.kv_offload.base import (
    ReqContext,
    make_offload_key,
 )
-from vllm.v1.kv_offload.tiering.base import JobMetadata
-from vllm.v1.kv_offload.tiering.example import ExampleSecondaryTier
+from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
 from vllm.v1.kv_offload.tiering.manager import (
    CPUPrimaryTierOffloadingManager,
    TieringOffloadingManager,
@@ -63,14 +61,17 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None:
    return count


-class TestExampleSecondaryTier:
-    """Tests for ExampleSecondaryTier implementation."""
+class TestExampleSecondaryTierManager:
+    """Tests for ExampleSecondaryTierManager implementation."""

    def test_basic_store_and_lookup(self):
        """Test basic store and lookup operations."""
        mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
-        tier = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
+        tier = ExampleSecondaryTierManager(
+            vllm_config=_MOCK_VLLM_CONFIG,
+            primary_kv_view=mock_view,
+            tier_type="example",
+            custom_param=67,
        )

        # Initially empty
@@ -88,81 +89,6 @@ class TestExampleSecondaryTier:
        # Third block not present
        assert tier.lookup(blocks[2], _CTX) is False

-    def test_lru_eviction(self):
-        """Test LRU eviction policy."""
-        mock_view = memoryview(torch.zeros((4, 16), dtype=torch.int8).numpy())
-        tier = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=3
-        )
-
-        # Fill tier to capacity
-        blocks = to_keys(range(3))
-        for block in blocks:
-            tier.blocks[block] = True
-
-        assert tier.get_num_blocks() == 3
-
-        # Touch first block (make it most recently used)
-        tier.touch([blocks[0]], _CTX)
-
-        # Store new block should evict blocks[1] (least recently used)
-        new_block = to_keys([3])[0]
-
-        tier.submit_store(
-            JobMetadata(
-                job_id=1,
-                keys=[new_block],
-                block_ids=np.array([0], dtype=np.int64),
-                is_promotion=False,
-                req_context=_CTX,
-            )
-        )
-
-        # Complete the job
-        tier.get_finished()
-
-        # Verify new block is stored and blocks[1] was evicted (LRU)
-        assert new_block in tier.blocks
-        assert blocks[1] not in tier.blocks
-        # blocks[0] and blocks[2] should still be present
-        assert blocks[0] in tier.blocks
-        assert blocks[2] in tier.blocks
-
-    def test_async_simulation(self):
-        """Test simulated async behavior."""
-        mock_view = memoryview(torch.zeros((10, 16), dtype=torch.int8).numpy())
-        tier = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG,
-            primary_kv_view=mock_view,
-            max_blocks=10,
-            simulate_async=True,
-        )
-
-        blocks = to_keys(range(2))
-
-        # Submit store job
-        tier.submit_store(
-            JobMetadata(
-                job_id=1,
-                keys=blocks,
-                block_ids=np.array([0, 1], dtype=np.int64),
-                is_promotion=False,
-                req_context=_CTX,
-            )
-        )
-
-        # Blocks should not yet be stored (pending async completion)
-        assert tier.get_num_blocks() == 0
-
-        # First get_finished() should complete the job
-        completed = list(tier.get_finished())
-        assert len(completed) == 1
-        assert completed[0].job_id == 1
-        assert completed[0].success is True
-
-        # Blocks should now be stored
-        assert tier.get_num_blocks() == 2
-

 class TestTieringOffloadingManager:
    """Tests for TieringOffloadingManager."""
@@ -178,11 +104,15 @@ class TestTieringOffloadingManager:
        mock_view = mock_region.create_kv_memoryview()

        # Create secondary tiers with the primary view
-        self.secondary_tier1 = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
+        self.secondary_tier1 = ExampleSecondaryTierManager(
+            vllm_config=_MOCK_VLLM_CONFIG,
+            primary_kv_view=mock_view,
+            tier_type="example",
        )
-        self.secondary_tier2 = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG, primary_kv_view=mock_view, max_blocks=10
+        self.secondary_tier2 = ExampleSecondaryTierManager(
+            vllm_config=_MOCK_VLLM_CONFIG,
+            primary_kv_view=mock_view,
+            tier_type="example",
        )

        # Create tiered manager
@@ -359,20 +289,19 @@ class TestTieringOffloadingManager:
        self.manager.complete_store(blocks, _CTX, success=True)
        list(self.manager.take_events())

+        self.secondary_tier1.touch = MagicMock(wraps=self.secondary_tier1.touch)
+        self.secondary_tier2.touch = MagicMock(wraps=self.secondary_tier2.touch)
+
        # Touch blocks
        self.manager.touch(blocks, _CTX)

        # Verify touch was called on primary tier (check LRU order)
-        # In LRU, touched blocks should be at the end
        primary_keys = list(self.primary_tier._policy.blocks.keys())
        assert primary_keys[-3:] == list(reversed(blocks))

-        # Verify touch was called on all secondary tiers
-        secondary1_keys = list(self.secondary_tier1.blocks.keys())
-        assert secondary1_keys[-3:] == list(reversed(blocks))
-
-        secondary2_keys = list(self.secondary_tier2.blocks.keys())
-        assert secondary2_keys[-3:] == list(reversed(blocks))
+        # Verify touch was propagated to all secondary tiers
+        self.secondary_tier1.touch.assert_called_once_with(blocks, _CTX)
+        self.secondary_tier2.touch.assert_called_once_with(blocks, _CTX)

    def test_failed_store_no_cascade(self, manager_setup):
        """Test that failed GPU→primary store doesn't cascade."""
@@ -396,59 +325,6 @@ class TestTieringOffloadingManager:
        self.secondary_tier1.submit_store.assert_not_called()
        self.secondary_tier2.submit_store.assert_not_called()

-    def test_multiple_secondary_tiers_independent_eviction(self):
-        """Test that secondary tiers manage their own evictions."""
-        mock_region = _mock_mmap_region(10)
-        mock_view = mock_region.create_kv_memoryview()
-
-        # Create tier with small capacity
-        small_tier = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG,
-            primary_kv_view=mock_view,
-            max_blocks=5,
-            simulate_async=False,
-        )
-        large_tier = ExampleSecondaryTier(
-            vllm_config=_MOCK_VLLM_CONFIG,
-            primary_kv_view=mock_view,
-            max_blocks=10,
-            simulate_async=False,
-        )
-
-        # Create a fresh primary tier for this test
-        primary_tier = CPUPrimaryTierOffloadingManager(
-            num_blocks=10, mmap_region=mock_region
-        )
-
-        manager = TieringOffloadingManager(
-            primary_tier=primary_tier,
-            secondary_tiers=[small_tier, large_tier],
-        )
-
-        # First, store 5 blocks to fill the small tier
-        blocks1 = to_keys(range(5))
-        result = manager.prepare_store(blocks1, _CTX)
-        assert result is not None
-        manager.complete_store(blocks1, _CTX, success=True)
-        list(manager.take_events())
-
-        # Both tiers should have 5 blocks
-        assert small_tier.get_num_blocks() == 5
-        assert large_tier.get_num_blocks() == 5
-
-        # Now store 3 more blocks - small tier should evict 3 blocks
-        blocks2 = to_keys(range(5, 8))
-        result = manager.prepare_store(blocks2, _CTX)
-        assert result is not None
-        manager.complete_store(blocks2, _CTX, success=True)
-        list(manager.take_events())
-
-        # Small tier should still have 5 blocks (evicted 3, added 3)
-        assert small_tier.get_num_blocks() == 5
-
-        # Large tier should have all 8 blocks
-        assert large_tier.get_num_blocks() == 8
-
    def test_lookup_batches_submit_load_per_request(self, manager_setup):
        """lookup() defers submit_load until take_events(), one call per request.

@@ -15,7 +15,12 @@ _TORCH_CUDA_PATTERNS = [
    r"\bcuda_device_count_stateless\(\)\b",
 ]

-ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}
+ALLOWED_FILES = {
+    "vllm/platforms/",
+    "vllm/device_allocator/",
+    "vllm/distributed/weight_transfer/ipc_engine.py",
+    "tests/distributed/test_packed_tensor.py",
+}


 def scan_file(path: str) -> int:
@@ -878,7 +878,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
        return {}

    # Analyze the functions to determine FA3/FA4-specific features
-    fa3_supports_fp8 = False
+    fa3_supports_fp8 = True
    fa3_supports_sinks = False
    fa4_supports_sinks = False
    fa3_compute_cap: str | None = None
@@ -888,18 +888,6 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
        if not isinstance(node, ast.FunctionDef):
            continue

-        # Check flash_attn_supports_fp8 - looks for `get_flash_attn_version() == 3`
-        if node.name == "flash_attn_supports_fp8":
-            for n in ast.walk(node):
-                if (
-                    isinstance(n, ast.Compare)
-                    and isinstance(n.left, ast.Call)
-                    and isinstance(n.left.func, ast.Name)
-                    and n.left.func.id == "get_flash_attn_version"
-                ):
-                    fa3_supports_fp8 = True
-                    break
-
        # Check flash_attn_supports_sinks - looks for `fa_version == 3/4`
        # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)
        if node.name == "flash_attn_supports_sinks":
@@ -414,17 +414,32 @@ _AITER_HAS_FUSED_QK_RMSNORM: bool | None = None


 def check_aiter_fused_qk_rmsnorm() -> bool:
-    """Check if aiter provides fused_qk_rmsnorm (requires AITer >= PR #2442)."""
+    """Check if aiter provides fused_qk_rmsnorm.
+
+    Supports both the new private name ``_fused_qk_rmsnorm``
+    (AITER >= PR #2958) and the old public name ``fused_qk_rmsnorm``
+    (AITER >= PR #2442).
+
+    TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
+    AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
+    """
    global _AITER_HAS_FUSED_QK_RMSNORM
    if _AITER_HAS_FUSED_QK_RMSNORM is None:
        try:
            from aiter.ops.fused_qk_norm_rope_cache_quant import (  # noqa: F401
-                fused_qk_rmsnorm,
+                _fused_qk_rmsnorm,
            )

            _AITER_HAS_FUSED_QK_RMSNORM = True
        except (ImportError, ModuleNotFoundError, AttributeError):
-            _AITER_HAS_FUSED_QK_RMSNORM = False
+            try:
+                from aiter.ops.fused_qk_norm_rope_cache_quant import (  # noqa: F401
+                    fused_qk_rmsnorm,
+                )
+
+                _AITER_HAS_FUSED_QK_RMSNORM = True
+            except (ImportError, ModuleNotFoundError, AttributeError):
+                _AITER_HAS_FUSED_QK_RMSNORM = False
    return _AITER_HAS_FUSED_QK_RMSNORM


@@ -744,7 +759,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
    total_bytes = input_.numel() * input_.element_size()
    hidden_dim = input_.shape[-1]
    token_num = input_.shape[0]
-    hidden_ok = hidden_dim in (512, 1024, 2048, 4096, 7168)
+    if input_.dtype in (torch.bfloat16, torch.float16):
+        pack_size = 16 // input_.element_size()
+        hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
+    else:
+        hidden_ok = False
    token_ok = token_num <= 80
    world_size = aiter_ar.world_size
    full_nvlink = aiter_ar.fully_connected
@@ -1062,21 +1081,42 @@ def _fused_mla_dual_rms_norm_impl(
    x2_epsilon: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    try:
-        from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
-    except (ImportError, ModuleNotFoundError) as exc:
+        import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops
+    except (ImportError, ModuleNotFoundError, AttributeError) as exc:
        raise ImportError(
-            "fused_qk_rmsnorm requires a newer AITer version "
-            "(>= PR #2442). Please upgrade aiter or disable the "
+            "fused_qk_rmsnorm requires AITer >= PR #2442. "
+            "Please upgrade aiter or disable the "
            "fuse_mla_dual_rms_norm pass."
        ) from exc

-    return fused_qk_rmsnorm(
-        q=x1,
-        q_weight=x1_weight,
-        q_eps=x1_epsilon,
-        k=x2,
-        k_weight=x2_weight,
-        k_eps=x2_epsilon,
+    if hasattr(aiter_ops, "_fused_qk_rmsnorm"):
+        return aiter_ops._fused_qk_rmsnorm(
+            q_out=None,
+            q=x1,
+            q_weight=x1_weight,
+            q_eps=x1_epsilon,
+            k_out=None,
+            k=x2,
+            k_weight=x2_weight,
+            k_eps=x2_epsilon,
+        )
+
+    # TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
+    # AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
+    if hasattr(aiter_ops, "fused_qk_rmsnorm"):
+        return aiter_ops.fused_qk_rmsnorm(
+            q=x1,
+            q_weight=x1_weight,
+            q_eps=x1_epsilon,
+            k=x2,
+            k_weight=x2_weight,
+            k_eps=x2_epsilon,
+        )
+
+    raise ImportError(
+        "fused_qk_rmsnorm requires AITer >= PR #2442. "
+        "Please upgrade aiter or disable the "
+        "fuse_mla_dual_rms_norm pass."
    )


@@ -280,123 +280,6 @@ def merge_attn_states(
    )


-def convert_vertical_slash_indexes(
-    q_seqlens: torch.Tensor,  # [BATCH, ]
-    kv_seqlens: torch.Tensor,  # [BATCH, ]
-    vertical_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_V]
-    slash_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_S]
-    context_size: int,
-    block_size_M: int,
-    block_size_N: int,
-    causal: bool = True,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    batch_size = slash_indexes.size(0)
-    num_heads = slash_indexes.size(1)
-    nnz_slash = slash_indexes.size(2)
-    nnz_vertical = vertical_indexes.size(2)
-    num_rows = (context_size + block_size_M - 1) // block_size_M
-
-    block_count = torch.zeros(
-        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
-    )
-    block_offset = torch.zeros(
-        batch_size,
-        num_heads,
-        num_rows,
-        nnz_slash,
-        dtype=q_seqlens.dtype,
-        device=q_seqlens.device,
-    )
-    column_count = torch.zeros(
-        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
-    )
-    column_index = torch.zeros(
-        batch_size,
-        num_heads,
-        num_rows,
-        nnz_vertical,
-        dtype=q_seqlens.dtype,
-        device=q_seqlens.device,
-    )
-
-    torch.ops._C.convert_vertical_slash_indexes(
-        block_count,
-        block_offset,
-        column_count,
-        column_index,
-        q_seqlens,
-        kv_seqlens,
-        vertical_indexes,
-        slash_indexes,
-        context_size,
-        block_size_M,
-        block_size_N,
-        causal,
-    )
-    return block_count, block_offset, column_count, column_index
-
-
-def convert_vertical_slash_indexes_mergehead(
-    q_seqlens: torch.Tensor,  # [BATCH, ]
-    kv_seqlens: torch.Tensor,  # [BATCH, ]
-    vertical_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_V]
-    slash_indexes: torch.Tensor,  # [BATCH, N_HEADS, NNZ_S]
-    # [N_HEADS] : different head use different number of indices
-    vertical_indices_count: torch.Tensor,
-    slash_indices_count: torch.Tensor,
-    context_size: int,
-    block_size_M: int,
-    block_size_N: int,
-    causal: bool = True,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    batch_size = slash_indexes.size(0)
-    num_heads = slash_indexes.size(1)
-    nnz_slash = slash_indexes.size(2)
-    nnz_vertical = vertical_indexes.size(2)
-    num_rows = (context_size + block_size_M - 1) // block_size_M
-
-    block_count = torch.empty(
-        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
-    )
-    block_offset = torch.empty(
-        batch_size,
-        num_heads,
-        num_rows,
-        nnz_slash,
-        dtype=q_seqlens.dtype,
-        device=q_seqlens.device,
-    )
-    column_count = torch.empty(
-        batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
-    )
-    column_index = torch.empty(
-        batch_size,
-        num_heads,
-        num_rows,
-        nnz_vertical,
-        dtype=q_seqlens.dtype,
-        device=q_seqlens.device,
-    )
-
-    torch.ops._C.convert_vertical_slash_indexes_mergehead(
-        block_count,
-        block_offset,
-        column_count,
-        column_index,
-        q_seqlens,
-        kv_seqlens,
-        vertical_indexes,
-        slash_indexes,
-        vertical_indices_count,
-        slash_indices_count,
-        context_size,
-        block_size_M,
-        block_size_N,
-        causal,
-    )
-    return block_count, block_offset, column_count, column_index
-
-
 # pos encoding ops
 def rotary_embedding(
    positions: torch.Tensor,
@@ -2658,32 +2541,7 @@ def moe_wna16_marlin_gemm(
    )


-if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
-
-    @register_fake("_moe_C::marlin_gemm_moe")
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"):

    @register_fake("_moe_C::moe_wna16_marlin_gemm")
    def moe_wna16_marlin_gemm_fake(
@@ -3195,6 +3053,7 @@ class CPUQuantMethod(IntEnum):
    INT8_W8A8 = 1
    FP8_W8A16 = 2
    INT4_W4A8 = 3
+    MXFP4 = 4


 if hasattr(torch.ops._C, "fused_experts_cpu"):
@@ -3213,6 +3072,10 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):
        w1_zero: torch.Tensor | None,
        w2_zero: torch.Tensor | None,
        block_size: list[int] | None,
+        w1_bias: torch.Tensor | None,
+        w2_bias: torch.Tensor | None,
+        alpha: float | None,
+        limit: float | None,
        is_vnni: bool,
    ) -> torch.Tensor:
        return torch.empty_like(hidden_states)
@@ -3231,7 +3094,11 @@ def fused_experts_cpu(
    w1_zero: torch.Tensor | None,
    w2_zero: torch.Tensor | None,
    block_size: list[int] | None,
-    is_vnni: bool,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    alpha: float | None = None,
+    limit: float | None = None,
+    is_vnni: bool = True,
 ) -> torch.Tensor:
    return torch.ops._C.fused_experts_cpu(
        hidden_states,
@@ -3246,6 +3113,10 @@ def fused_experts_cpu(
        w1_zero,
        w2_zero,
        block_size,
+        w1_bias,
+        w2_bias,
+        alpha,
+        limit,
        is_vnni,
    )

@@ -3370,6 +3241,135 @@ def fp8_scaled_mm_cpu(
    )


+def chunk_gated_delta_rule_cpu(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.Tensor,
+    head_first: bool,
+    use_qk_l2norm_in_kernel: bool,
+    eps: float = 1e-5,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.ops._C.chunk_gated_delta_rule_cpu(
+        query,
+        key,
+        value,
+        g,
+        beta,
+        initial_state,
+        output_final_state,
+        cu_seqlens,
+        head_first,
+        use_qk_l2norm_in_kernel,
+        eps,
+    )
+
+
+def fused_sigmoid_gating_delta_rule_update_cpu(
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    initial_state_source: torch.Tensor,
+    initial_state_indices: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    use_qk_l2norm_in_kernel: bool,
+    softplus_beta: float = 1.0,
+    softplus_threshold: float = 20.0,
+) -> torch.Tensor:
+    return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        initial_state_source,
+        initial_state_indices,
+        cu_seqlens,
+        use_qk_l2norm_in_kernel,
+        softplus_beta,
+        softplus_threshold,
+    )
+
+
+def fused_gdn_gating_cpu(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    dt_bias: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.ops._C.fused_gdn_gating_cpu(
+        A_log,
+        a,
+        b,
+        dt_bias,
+    )
+
+
+def causal_conv1d_weight_pack(
+    weight: torch.Tensor,
+) -> torch.Tensor:
+    return torch.ops._C.causal_conv1d_weight_pack(
+        weight,
+    )
+
+
+def causal_conv1d_fwd_cpu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    conv_states: torch.Tensor | None,
+    query_start_loc: torch.Tensor | None,
+    cache_indices: torch.Tensor | None,
+    has_initial_state: torch.Tensor | None,
+    silu_activation: bool,
+    is_vnni: bool,
+) -> torch.Tensor:
+    return torch.ops._C.causal_conv1d_fwd_cpu(
+        x,
+        weight,
+        bias,
+        conv_states,
+        query_start_loc,
+        cache_indices,
+        has_initial_state,
+        silu_activation,
+        -1,
+        is_vnni,
+    )
+
+
+def causal_conv1d_update_cpu(
+    x: torch.Tensor,
+    conv_states: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    silu_activation: bool,
+    conv_state_indices: torch.Tensor | None,
+    is_vnni: bool,
+) -> torch.Tensor:
+    return torch.ops._C.causal_conv1d_update_cpu(
+        x,
+        conv_states,
+        weight,
+        bias,
+        silu_activation,
+        None,
+        conv_state_indices,
+        -1,
+        is_vnni,
+    )
+
+
 class CPUDNNLGEMMHandler:
    def __init__(self) -> None:
        self.handler_tensor: torch.Tensor | None = None
@@ -1940,6 +1940,19 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
        ):
            dataset_class = MTBenchDataset
            args.hf_split = args.hf_split if args.hf_split else "train"
+        elif (
+            args.dataset_path in HumanEvalDataset.SUPPORTED_DATASET_PATHS
+            or args.hf_name in HumanEvalDataset.SUPPORTED_DATASET_PATHS
+        ):
+            dataset_class = HumanEvalDataset
+            args.hf_split = args.hf_split if args.hf_split else "test"
+        elif (
+            args.dataset_path in GSM8KDataset.SUPPORTED_DATASET_PATHS
+            or args.hf_name in GSM8KDataset.SUPPORTED_DATASET_PATHS
+        ):
+            dataset_class = GSM8KDataset
+            args.hf_subset = args.hf_subset if args.hf_subset else "main"
+            args.hf_split = args.hf_split if args.hf_split else "test"
        elif (
            args.dataset_path in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
            or args.hf_name in MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
@@ -3122,6 +3135,126 @@ class MTBenchDataset(HuggingFaceDataset):
        return sampled_requests


+# -----------------------------------------------------------------------------
+# HumanEval Dataset Implementation
+# -----------------------------------------------------------------------------
+
+
+class HumanEvalDataset(HuggingFaceDataset):
+    """
+    HumanEvalDataset Dataset.
+    https://huggingface.co/datasets/openai/openai_humaneval
+
+    We create a single turn dataset for HumanEval.
+    """
+
+    DEFAULT_OUTPUT_LEN = 256
+    SUPPORTED_DATASET_PATHS = {
+        "openai/openai_humaneval",
+    }
+
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        skip_chat_template: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]:
+        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
+        sampled_requests = []
+
+        for i, item in enumerate(self.data):
+            if len(sampled_requests) >= num_requests:
+                break
+            prompt = item["prompt"]
+
+            # apply template
+            if not skip_chat_template:
+                prompt = tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+
+            prompt_len = len(tokenizer(prompt).input_ids)
+            sampled_requests.append(
+                SampleRequest(
+                    prompt=prompt,
+                    prompt_len=prompt_len,
+                    expected_output_len=output_len,
+                    request_id=request_id_prefix + str(i),
+                )
+            )
+        self.maybe_oversample_requests(
+            sampled_requests, num_requests, request_id_prefix, no_oversample
+        )
+        return sampled_requests
+
+
+# -----------------------------------------------------------------------------
+# GSM8K Dataset Implementation
+# -----------------------------------------------------------------------------
+
+
+class GSM8KDataset(HuggingFaceDataset):
+    """
+    GSM8K Dataset.
+    https://huggingface.co/datasets/openai/gsm8k
+
+    We create a single turn dataset for GSM8K.
+    """
+
+    DEFAULT_OUTPUT_LEN = 256
+    SUPPORTED_DATASET_PATHS = {
+        "openai/gsm8k",
+    }
+
+    def sample(
+        self,
+        tokenizer: TokenizerLike,
+        num_requests: int,
+        request_id_prefix: str = "",
+        no_oversample: bool = False,
+        output_len: int | None = None,
+        enable_multimodal_chat: bool = False,
+        skip_chat_template: bool = False,
+        **kwargs,
+    ) -> list[SampleRequest]:
+        output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN
+        sampled_requests = []
+
+        for i, item in enumerate(self.data):
+            if len(sampled_requests) >= num_requests:
+                break
+            prompt = item["question"]
+
+            # apply template
+            if not skip_chat_template:
+                prompt = tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+
+            prompt_len = len(tokenizer(prompt).input_ids)
+            sampled_requests.append(
+                SampleRequest(
+                    prompt=prompt,
+                    prompt_len=prompt_len,
+                    expected_output_len=output_len,
+                    request_id=request_id_prefix + str(i),
+                )
+            )
+        self.maybe_oversample_requests(
+            sampled_requests, num_requests, request_id_prefix, no_oversample
+        )
+        return sampled_requests
+
+
 # -----------------------------------------------------------------------------
 # Blazedit Dataset Implementation
 # -----------------------------------------------------------------------------
@@ -0,0 +1,424 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Breakable CUDA graph capture/replay.
+
+This is an alternative to :class:`CUDAGraphWrapper` that replaces vLLM's
+torch.compile-based FX graph splitting with runtime stream-capture
+breaks.
+
+The idea (inspired by sgl-project/sglang#19102): instead of pre-splitting
+the model into many pieces at attention boundaries, a
+single capture context drives the whole forward and intercepts
+attention / kv-cache custom ops at the dispatcher to end the current
+stream capture, run the op eagerly, and resume capture.
+
+The captured artifact is a list of zero-arg callables -- the bound
+``CUDAGraph.replay`` for graph segments, or the user fn for eager
+segments -- replayed in order at inference time.
+
+Eager segments must operate on the same static buffers used during
+capture so subsequent graph segments read the same memory addresses.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+import gc
+import threading
+import weakref
+from collections.abc import Callable
+from typing import Any, ClassVar, TypeVar
+
+import torch
+
+import vllm.envs as envs
+from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
+from vllm.forward_context import (
+    BatchDescriptor,
+    get_forward_context,
+    is_forward_context_available,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.offloader.base import get_offloader
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import weak_ref_tensor, weak_ref_tensors
+
+logger = init_logger(__name__)
+
+
+def is_breakable_cudagraph_enabled() -> bool:
+    return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def eager_break_during_capture(fn: F) -> F:
+    """Decorator that turns a custom-op Python kernel into a "break point"
+    for the breakable cudagraph capture.
+
+    When the decorated function is invoked outside of a
+    :class:`BreakableCUDAGraphCapture` context, it executes normally.
+
+    When invoked inside a capture context, it ends the current cudagraph
+    segment, runs the function eagerly on the capture stream, records the
+    callable for replay, and starts a fresh segment.
+
+    **In-place output buffer required.** Decorated ops must write into a
+    caller-provided output tensor; a fresh tensor returned by ``fn`` would
+    change address each replay and break downstream graph segments.
+
+    **Decorator order matters.** Apply as the *outermost* decorator if
+    there are other decorators that introduce host-side side effects
+    around the call -- the canonical example is
+    ``@maybe_transfer_kv_layer`` for PD-disaggregation, whose
+    ``wait_for_layer_load`` and ``save_kv_layer`` calls must run in the
+    eager segment, not inside the captured cudagraph. Putting
+    ``@eager_break_during_capture`` *inside* such a decorator would
+    record those side effects into the graph and hang on replay.
+
+    The correct order is::
+
+        @eager_break_during_capture   # outermost
+        @maybe_transfer_kv_layer
+        def unified_attention_with_output(...):
+            ...
+    """
+    if not is_breakable_cudagraph_enabled():
+        return fn
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        capture = BreakableCUDAGraphCapture.current()
+        if capture is None:
+            return fn(*args, **kwargs)
+        if not capture._capturing:
+            return fn(*args, **kwargs)
+        if is_forward_context_available():
+            mode = get_forward_context().cudagraph_runtime_mode
+            if mode == CUDAGraphMode.FULL:
+                return fn(*args, **kwargs)
+
+        # Weak-ref args: strong refs in the replay lambda pin cudagraph-pool
+        # slots across batch descriptors. cudagraph owns the slot, so the
+        # weak_ref is safe to deref on replay.
+        weak_args = tuple(
+            weak_ref_tensor(a) if isinstance(a, torch.Tensor) else a for a in args
+        )
+        weak_kwargs = {
+            k: weak_ref_tensor(v) if isinstance(v, torch.Tensor) else v
+            for k, v in kwargs.items()
+        }
+        return capture.add_eager(lambda: fn(*weak_args, **weak_kwargs))
+
+    return wrapper  # type: ignore[return-value]
+
+
+# ---------------------------------------------------------------------------
+# Capture context
+# ---------------------------------------------------------------------------
+
+
+class BreakableCUDAGraphCapture:
+    """Stream-capture context that supports eager breaks via :meth:`add_eager`.
+
+    Usage::
+
+        cap = BreakableCUDAGraphCapture(pool=...)
+        with cap:
+            output = model(*static_inputs)
+        # Later, after copying new inputs into the static buffers:
+        cap.replay()
+        # Output tensors live at the same addresses as during capture.
+
+    Thread-local: only one capture may be active per thread.
+    """
+
+    _tls = threading.local()
+
+    @classmethod
+    def current(cls) -> BreakableCUDAGraphCapture | None:
+        return getattr(cls._tls, "active", None)
+
+    @classmethod
+    def is_active(cls) -> bool:
+        return cls.current() is not None
+
+    def __init__(self, pool: Any | None = None) -> None:
+        self.pool = pool
+        self.segments: list[Callable[[], Any]] = []
+        self._num_graphs: int = 0
+        self._num_eager_breaks: int = 0
+        self._current_graph: torch.cuda.CUDAGraph | None = None
+        self._capturing: bool = False
+
+    # --- context manager protocol ----------------------------------------
+
+    def __enter__(self) -> BreakableCUDAGraphCapture:
+        if getattr(BreakableCUDAGraphCapture._tls, "active", None) is not None:
+            raise RuntimeError("Nested BreakableCUDAGraphCapture is not supported.")
+        BreakableCUDAGraphCapture._tls.active = self
+        self._begin_segment()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        try:
+            self._end_segment()
+        finally:
+            BreakableCUDAGraphCapture._tls.active = None
+
+    # --- segment management ----------------------------------------------
+
+    def _begin_segment(self) -> None:
+        assert not self._capturing
+        g = torch.cuda.CUDAGraph()
+        if self.pool is not None:
+            g.capture_begin(pool=self.pool)
+        else:
+            g.capture_begin()
+        self._current_graph = g
+        self._capturing = True
+
+    def _end_segment(self) -> None:
+        if not self._capturing:
+            return
+        assert self._current_graph is not None
+        self._current_graph.capture_end()
+        self.segments.append(self._current_graph.replay)
+        self._num_graphs += 1
+        self._current_graph = None
+        self._capturing = False
+
+    def add_eager(self, fn: Callable[[], Any]) -> Any:
+        """End the current capture segment, run ``fn`` eagerly on the
+        capture stream, record ``fn`` for replay, and start a new segment.
+
+        Returns whatever ``fn`` returned during this (capture-time) call.
+        Replay does not return values; callers should propagate any
+        downstream dependencies via static output buffers.
+        """
+        self._end_segment()
+        result = fn()
+        self.segments.append(fn)
+        self._num_eager_breaks += 1
+        self._begin_segment()
+        return result
+
+    # --- replay ----------------------------------------------------------
+
+    def replay(self) -> None:
+        for r in self.segments:
+            r()
+
+    # --- introspection ---------------------------------------------------
+
+    @property
+    def num_graphs(self) -> int:
+        return self._num_graphs
+
+    @property
+    def num_eager_breaks(self) -> int:
+        return self._num_eager_breaks
+
+    def __repr__(self) -> str:
+        return (
+            f"BreakableCUDAGraphCapture(graphs={self.num_graphs}, "
+            f"eager_breaks={self.num_eager_breaks})"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Wrapper that mirrors CUDAGraphWrapper's interface
+# ---------------------------------------------------------------------------
+
+
+@dataclasses.dataclass
+class _BreakableEntry:
+    batch_descriptor: BatchDescriptor
+    capture: BreakableCUDAGraphCapture | None = None
+    output: Any = None
+    input_addresses: list[int] | None = None
+
+
+class BreakableCUDAGraphWrapper:
+    """Drop-in replacement for :class:`CUDAGraphWrapper` that uses
+    :class:`BreakableCUDAGraphCapture` instead of a single monolithic
+    ``torch.cuda.graph()`` capture.
+
+    Same dispatch contract as ``CUDAGraphWrapper``:
+        * If no ``forward_context`` is available, run the underlying
+          callable eagerly.
+        * If runtime mode mismatch / NONE, run eagerly.
+        * Otherwise, lazily capture per ``batch_descriptor`` and replay
+          on subsequent invocations with the same descriptor.
+    """
+
+    _all_instances: ClassVar[weakref.WeakSet[BreakableCUDAGraphWrapper]] = (
+        weakref.WeakSet()
+    )
+
+    @classmethod
+    def clear_all_graphs(cls) -> None:
+        for instance in list(cls._all_instances):
+            instance.clear_graphs()
+
+    def __init__(
+        self,
+        runnable: Callable[..., Any],
+        vllm_config: VllmConfig,
+    ) -> None:
+        # Unlike the original CUDAGraphWrapper which strictly matches a
+        # single runtime_mode, this wrapper captures whatever the
+        # dispatcher emits (any non-NONE runtime_mode) -- breakable's
+        # capture is identical for prefill and decode, so there's nothing
+        # to dispatch on at the runtime_mode level. Entries are keyed by
+        # BatchDescriptor which already encodes batch shape / uniformity.
+        self.runnable = runnable
+        self.vllm_config = vllm_config
+        self.compilation_config = vllm_config.compilation_config
+        self.graph_pool = current_platform.get_global_graph_pool()
+        self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
+
+        self.entries: dict[BatchDescriptor, _BreakableEntry] = {}
+        BreakableCUDAGraphWrapper._all_instances.add(self)
+
+        logger.info_once("Breakable CUDA graph enabled")
+
+    # --- vllm-style attribute forwarding ---------------------------------
+
+    def __getattr__(self, key: str) -> Any:
+        runnable = self.__dict__.get("runnable")
+        if runnable is not None and hasattr(runnable, key):
+            return getattr(runnable, key)
+        raise AttributeError(key)
+
+    def unwrap(self) -> Callable[..., Any]:
+        return self.runnable
+
+    @property
+    def cudagraph_wrapper(self) -> BreakableCUDAGraphWrapper:
+        return self
+
+    def clear_graphs(self) -> None:
+        self.entries.clear()
+
+    # --- dispatch --------------------------------------------------------
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if not is_forward_context_available():
+            return self.runnable(*args, **kwargs)
+
+        forward_context = get_forward_context()
+        batch_descriptor = forward_context.batch_descriptor
+        cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
+
+        # Capture whenever the dispatcher says "some cudagraph mode" --
+        # breakable produces the same artifact regardless of PIECEWISE
+        # vs FULL, so we match either. Entries are keyed by batch
+        # descriptor, which already encodes prefill/decode distinctions.
+        if cudagraph_runtime_mode == CUDAGraphMode.NONE:
+            return self.runnable(*args, **kwargs)
+
+        assert batch_descriptor is not None
+        entry = self.entries.get(batch_descriptor)
+        if entry is None:
+            entry = _BreakableEntry(batch_descriptor=batch_descriptor)
+            self.entries[batch_descriptor] = entry
+
+        if entry.capture is None:
+            return self._capture(entry, args, kwargs)
+        return self._replay(entry, args, kwargs)
+
+    # --- capture / replay paths -----------------------------------------
+
+    @staticmethod
+    def _collect_tensor_addresses(
+        args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> list[int]:
+        """Flatten tensor data_ptrs from positional and keyword args in a
+        stable order (positionals first, then kwargs in insertion order).
+
+        Used for the DEBUG-mode address-stability check; covers both call
+        styles since vLLM models are typically invoked with kwargs.
+        """
+        addrs = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
+        addrs.extend(
+            v.data_ptr() for v in kwargs.values() if isinstance(v, torch.Tensor)
+        )
+        return addrs
+
+    def _capture(
+        self,
+        entry: _BreakableEntry,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        validate_cudagraph_capturing_enabled()
+
+        entry.input_addresses = self._collect_tensor_addresses(args, kwargs)
+
+        if self.graph_pool is not None:
+            set_graph_pool_id(self.graph_pool)
+        else:
+            set_graph_pool_id(current_platform.graph_pool_handle())
+
+        # Match torch.cuda.graph()'s pre-capture cleanup once per descriptor.
+        # We drive capture_begin/end directly and bypass torch.cuda.graph(),
+        # so its built-in gc + empty_cache never fire. Run them here once
+        # per _capture call -- NOT inside _begin_segment, since this capture
+        # session may issue many begin/end pairs (one per layer's break),
+        # and repeated gc would tank capture time the way it did for the
+        # pre-`gc_disable` piecewise path.
+        gc.collect()
+        torch.accelerator.empty_cache()
+        # Sync the offloader's copy stream before capture so any in-flight
+        # pre-capture prefetches are complete and don't leak into the graph.
+        get_offloader().sync_prev_onload()
+
+        capture = BreakableCUDAGraphCapture(pool=self.graph_pool)
+        with capture:
+            output = self.runnable(*args, **kwargs)
+            # Join the offloader's copy stream while we still hold the last
+            # segment open, so the join is captured into the graph (otherwise
+            # we get an "unjoined stream" error on subsequent forwards).
+            get_offloader().join_after_forward()
+            # Convert output to a weak ref *inside* the capture context so the
+            # strong ref is dropped before the last segment closes, letting
+            # the cudagraph pool reclaim/reuse that memory immediately for
+            # the next batch descriptor's capture.
+            output = weak_ref_tensors(output)
+
+        entry.capture = capture
+        entry.output = weak_ref_tensors(output)
+
+        logger.debug(
+            "Captured breakable cudagraph for %s: %r",
+            entry.batch_descriptor,
+            capture,
+        )
+        # Return the (already-weak) output from the captured run so the
+        # caller of model(...) gets a tensor pointing at the cudagraph pool's memory
+        return output
+
+    def _replay(
+        self,
+        entry: _BreakableEntry,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        if self.is_debugging_mode and entry.input_addresses is not None:
+            new_addresses = self._collect_tensor_addresses(args, kwargs)
+            assert new_addresses == entry.input_addresses, (
+                "Input tensor addresses changed between capture and replay "
+                f"for {entry.batch_descriptor}. Expected "
+                f"{entry.input_addresses}, got {new_addresses}."
+            )
+        # Sync the offloader's copy stream before replay so any external
+        # dependencies from pre-capture prefetches are satisfied.
+        get_offloader().sync_prev_onload()
+        assert entry.capture is not None
+        entry.capture.replay()
+        return entry.output
@@ -7,7 +7,7 @@ from typing import Any, ParamSpec, TypeVar
 from torch import fx as fx

 from vllm import envs
-from vllm._aiter_ops import rocm_aiter_ops
+from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops
 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.logger import init_logger
@@ -169,7 +169,11 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
                if rocm_aiter_ops.is_enabled():
                    self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]

-            if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
+            if (
+                self.pass_config.fuse_mla_dual_rms_norm
+                and rocm_aiter_ops.is_enabled()
+                and check_aiter_fused_qk_rmsnorm()
+            ):
                self.passes += [MLADualRMSNormFusionPass(config)]

            if self.pass_config.fuse_rope_kvcache:
@@ -32,7 +32,7 @@ CacheDType = Literal[
    "fp8_per_token_head",
    "nvfp4",
 ]
-MambaDType = Literal["auto", "float32", "float16"]
+MambaDType = Literal["auto", "float32", "float16", "bfloat16"]
 MambaCacheMode = Literal["all", "align", "none"]
 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
 KVOffloadingBackend = Literal["native", "lmcache"]
@@ -20,8 +20,8 @@ logger = init_logger(__name__)
 class IrOpPriorityConfig:
    """
    Configuration for vLLM IR op priority for dispatching/lowering during the
-    forward pass. Each member is a list of strings, which will be passed to
-    vllm.ir.ops.<op_name>.set_priority() for the duration of the forward pass.
+    forward pass. Each member is a list of strings, which will be installed
+    in worker init via vllm.ir.ops.<op_name>.set_default().
    A single comma-separated string is accepted as well,

    If specified manually, platform defaults will be appended to the lists.
@@ -67,6 +67,31 @@ class IrOpPriorityConfig:
        assert all(isinstance(v, str) for v in value)
        return value

+    def _iter_op_priorities(self):
+        """
+        Yield (IrOp, priority_list) for each field, after importing platform
+        kernels and validating each entry.
+        """
+        from vllm.ir.op import IrOp
+        from vllm.platforms import current_platform
+
+        current_platform.import_ir_kernels()
+
+        for field in fields(self):  # type: ignore[arg-type]
+            op_priority = getattr(self, field.name)
+            assert op_priority is not None, (
+                f"IR op priority for {field.name} must be set"
+            )
+            logger.debug("Setting IR op priority for %s to %s", field.name, op_priority)
+            yield IrOp.registry[field.name], op_priority
+
+    def set_default(self) -> None:
+        """
+        Permanently set the IR op priority for all op members.
+        """
+        for ir_op, op_priority in self._iter_op_priorities():
+            ir_op.set_default(op_priority)
+
    @contextlib.contextmanager
    def set_priority(self):
        """
@@ -74,23 +99,9 @@ class IrOpPriorityConfig:
        It also imports IR kernel implementations for the current platform
        to ensure all implementations are made available.
        """
-        from vllm.ir.op import IrOp
-        from vllm.platforms import current_platform
-
-        current_platform.import_ir_kernels()
-
        with contextlib.ExitStack() as stack:
-            for field in fields(self):  # type: ignore[arg-type]
-                op_priority = getattr(self, field.name)
-                assert op_priority is not None, (
-                    f"IR op priority for {field.name} must be set"
-                )
-                logger.debug(
-                    "Setting IR op priority for %s to %s", field.name, op_priority
-                )
-                ir_op = IrOp.registry[field.name]
+            for ir_op, op_priority in self._iter_op_priorities():
                stack.enter_context(ir_op.set_priority(op_priority))
-
            yield

    @classmethod
@@ -124,6 +135,24 @@ MoEBackend = Literal[
    "emulation",
 ]

+LinearBackend = Literal[
+    "auto",
+    "cutlass",
+    "flashinfer_cutlass",
+    "flashinfer_trtllm",
+    "flashinfer_cudnn",
+    "marlin",
+    "triton",
+    "deep_gemm",
+    "torch",
+    "aiter",
+    "machete",
+    "fbgemm",
+    "conch",
+    "exllama",
+    "emulation",
+]
+

@config
 class KernelConfig:
@@ -157,6 +186,25 @@ class KernelConfig:
                   running QDQ on activations.
    """

+    linear_backend: LinearBackend = "auto"
+    """Backend for quantized linear layer GEMM kernels. Available options:
+
+    - "auto": Automatically select the best backend based on model and hardware
+    - "cutlass": Use CUTLASS-based kernels
+    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
+    - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
+    - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
+    - "marlin": Use Marlin kernels
+    - "triton": Use Triton-based kernels
+    - "deep_gemm": Use DeepGEMM kernels
+    - "torch": Use PyTorch native scaled_mm kernels
+    - "aiter": Use AMD AITer kernels (ROCm only)
+    - "machete": Use Machete kernels (mixed-precision)
+    - "fbgemm": Use FBGEMM kernels
+    - "conch": Use Conch mixed-precision kernels
+    - "exllama": Use Exllama mixed-precision kernels
+    - "emulation": Use slow dequant-to-BF16 emulation (for testing only)"""
+
    @field_validator("moe_backend", mode="before")
    @classmethod
    def _normalize_moe_backend(cls, value: Any) -> Any:
@@ -164,6 +212,13 @@ class KernelConfig:
            return value.lower().replace("-", "_")
        return value

+    @field_validator("linear_backend", mode="before")
+    @classmethod
+    def _normalize_linear_backend(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            return value.lower().replace("-", "_")
+        return value
+
    def compute_hash(self) -> str:
        """
        Produces a hash unique to the pass configuration.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
yewentao256	c2de222aa3	Merge branch 'main' into wentao-optimize-model-runner-v2-sampler Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-18 22:45:59 +00:00
Wentao YeandGitHub	cd49a05d5a	[Refactor] Remove dead code (#42889 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-18 16:41:22 -04:00
Ronen SchafferandGitHub	84747489de	Tier offload followup (#42529 ) Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>	2026-05-18 19:41:58 +00:00
Tuukka SarviandGitHub	8fc1c284b9	[ROCm] Guard AITER GDN decode fast path by layout (#42880 ) Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>	2026-05-18 11:56:22 -07:00
Amit PortnoyandGitHub	ce88f01c9a	[Docs] update attribution to reflect EDEN foundation (#41666 ) Signed-off-by: amitport <1131991+amitport@users.noreply.github.com>	2026-05-18 11:22:56 -07:00
Wentao YeandGitHub	00e20e76f7	[Refactor] Remove dead cuda kernels (#42767 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-18 11:14:21 -07:00
czhu-cohereandGitHub	9758a6e5c5	[BugFix] support PP for Cohere vision model (#42819 ) Signed-off-by: <conway.zhu@cohere.com> Signed-off-by: root <conway.zhu@cohere.com>	2026-05-18 11:12:06 -07:00
Bowen BaoandGitHub	a2c8fc6657	[ROCm][Quantization][3/N] Refactor quark_moe w4a4 w/ oracle (#41436 ) Signed-off-by: Bowen Bao <bowenbao@amd.com>	2026-05-18 13:46:13 -04:00
Jinzhen Lin GitHub Michael Goin	6859ca7615	[Bugfix] fix swiglu limit issue for humming backend + deepseek v4 (#42541 ) Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-05-18 17:32:26 +00:00
Mohammad Miadh AngkadandGitHub	67f58ce23f	[Bugfix] Fix DSV4 MTP after ROCm mHC integration (#42930 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>	2026-05-18 17:02:01 +00:00
Wei ZhaoandGitHub	8c296de63b	[Perf] Re-enable flashinfer autotune by default and cleanup (#42857 ) Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>	2026-05-18 09:12:27 -07:00
Harry MellorandGitHub	b12745e4f3	Fix `--convert` passed without `--runner` on causal models (#42935 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-05-18 15:56:09 +00:00
Wentao YeandGitHub	e26736973a	[Model Runner V2] Fix prompt logprobs calculation `Sizes of tensors must match` error (#42778 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-18 08:27:21 -07:00
Netanel HaberandGitHub	47829b1159	[Bugfix] mamba: run single-token extends as decodes (#42430 ) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>	2026-05-18 15:26:00 +00:00
Blanc SwanandGitHub	4a39b4f553	[Model] Add Apertus Tool Parser (#41154 ) Signed-off-by: Blanc <swan.blanc@infomaniak.com>	2026-05-18 11:20:04 -04:00
Siddharth Bedekar GitHub Robert Shaw OpenAI Codexgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	78e7a7b9b0	Refactor AWQ Marlin MoE onto modular WNA16 oracle (#42483 ) Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com> Signed-off-by: Siddharth Bedekar <104613085+bedeks@users.noreply.github.com> Co-authored-by: Robert Shaw <robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-05-18 08:02:43 -07:00
Michael Goin GitHub Claude Opus 4.7	f5d3dc7115	[Model Runner v2] Support update_config (#42783 ) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-05-18 10:26:07 -04:00
vllm-agent GitHub Luka Govedič Zhewen Li	1ac10f159a	Revert "[torch.compile] Add patch for fullgraph compilation" (#42686 ) (#42913 ) Co-authored-by: Luka Govedič <luka.govedic@gmail.com> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>	2026-05-18 09:02:51 -04:00
liranschour GitHub Or Ozerimergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	e5417657e5	[KV Connector][Offloading] Flush all pending jobs on last step (#42611 ) Signed-off-by: Liran Schour <lirans@il.ibm.com> Signed-off-by: liranschour <liranschour@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-05-18 12:59:42 +00:00
xiangdongandGitHub	2e40faf08b	[XPU][CI] Temporarily skip test_moe_lora_align_block_size_mixed_base_and_lora[1] in Intel GPU CI (#42954 ) Signed-off-by: zengxian <xiangdong.zeng@intel.com>	2026-05-18 20:34:48 +08:00
Nicolò LucchesiandGitHub	69c91d010a	[MRv2] Default to MRv1 when a connector is present (#42955 ) Signed-off-by: NickLucche <nlucches@redhat.com>	2026-05-18 20:34:16 +08:00
roikoren755andGitHub	737bfa3a43	[Bugfix][Hybrid][NemotronH] Fix mamba_cache_mode=all + speculative decoding crash (#41233 ) Signed-off-by: Roi Koren <roik@nvidia.com>	2026-05-18 14:54:00 +03:00
Kfir ToledoandGitHub	e414e1f1c0	[Bugfix][KV Offload] count appended GPU blocks in store group_sizes (#42945 ) Signed-off-by: Kfir Toledo <kfir.toledo@ibm.com>	2026-05-18 11:36:02 +00:00
inisis GitHub Yongye Zhu	df852ed503	fix: remove unused norm for dpskv4 (#41710 ) Signed-off-by: inisis <desmond.yao@buaa.edu.cn> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-18 18:33:29 +08:00
Yuwen ZhouandGitHub	88a860d754	[CPU] Add MXFP4 W4A16 MoE support (#41922 ) Signed-off-by: yuwenzho <yuwen.zhou@intel.com> Signed-off-by: Yuwen Zhou <yuwen.zhou@intel.com>	2026-05-18 03:04:45 -07:00
Tianmu Li GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	cac81b6eda	[CPU Backend] Improve cpu thread utilization (#42666 ) Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-05-18 03:04:41 -07:00
Li, JiangandGitHub	b4601ad43f	[CPU] Add fused GDN support for AMX CPU platform (#42707 ) Signed-off-by: jiang1.li <jiang1.li@intel.com>	2026-05-18 03:04:36 -07:00
Jee Jee LiandGitHub	2267f70070	[Kernel] Pack topk id/weights triton kernel (#42527 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-05-18 03:04:31 -07:00
Tony Lin GitHubLi, Jiang <jiang1.li@intel.com>	965d076148	[CPU] Specify required KV cache layout for CPU attention backend (#42740 ) Signed-off-by: Tony Lin <tony.lin@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>	2026-05-18 17:38:54 +08:00
wenjun liu GitHub Kunshang Jimergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	c38bed4248	delete xpu ci (#42582 ) Signed-off-by: wenjun.liu <wenjun.liu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-05-18 16:36:45 +08:00
Xin YangandGitHub	998714b21b	[Perf] Add do_not_specialize in fused FP8 RoPE kernel (#42849 ) Signed-off-by: Xin Yang <xyangx@amazon.com>	2026-05-18 01:32:46 -07:00
Harry MellorandGitHub	9537542537	Revert checkpoint specific workaround in Transformers modelling backend (#42923 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-05-18 17:31:06 +09:00
Rishapveer SinghandGitHub	5ab6d1b3fd	[Model] [Perf] Use flatten for Qwen3.5's GDN output projection (#42311 ) Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.com>	2026-05-18 16:14:36 +08:00
Jee Jee Li GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>Roger Wang	7d5b033782	[LoRA] Support 2D and 3D MoE LoRA adapter at the same time (#42242 ) Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2026-05-18 15:22:26 +08:00
Nguyễn Thế Duy GitHub Jee Jee Li Jee Jee Li	e3aeee5ff8	[Bugfix] moe lora align kernel grid (#40131 ) Signed-off-by: TheDuyIT <nduy250299@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: dtnguyen <dtnguyen@nvidia.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>	2026-05-18 00:17:53 -07:00
Harry MellorandGitHub	c1f7854342	Improve logging when docs build is skipped (#42929 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-05-18 06:33:32 +00:00
gaozihao-shyandGitHub	23c15acd77	[BugFix] Kimi-K2.5: skip vision tower dtype conversion when using quantization (#42869 ) Signed-off-by: gaozihao-shy <gaozihao-shy@users.noreply.github.com> Signed-off-by: gaozihao <gaozihao3@huawei.com>	2026-05-18 05:07:16 +00:00
Andreas KaratzasandGitHub	b50646e5ef	[ROCm][CI] Stabilize ROCm pooling and multimodal CI (#42909 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-05-18 03:57:59 +00:00
SoyaazzandGitHub	990f49bdcb	[MM][CG] Enable encoder Cudagraph for Step3VL (#42224 ) Signed-off-by: JisoLya <523420504@qq.com> Signed-off-by: Soyaazz <523420504@qq.com>	2026-05-17 20:19:13 -07:00
Alec GitHub OpenAI Codex	107210442d	[CI] Add NIXL EP import canary (#42567 ) Signed-off-by: Alec Flowers <aflowers@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-05-17 19:11:46 -07:00
Yiliu Dong GitHub qianlihuang	03ddc1c9bc	[Perf] Wire silu_and_mul_per_block_quant into TritonFP8MoE (MiniMax-M2) (#42497 ) Signed-off-by: qianlihuang <yiliu.dong@qq.com> Signed-off-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: qianlihuang <yiliu.dong@qq.com>	2026-05-17 21:57:04 -04:00
Luka GovedičandGitHub	966903eb93	[torch.compile] Add patch for fullgraph compilation (#42686 ) Signed-off-by: Luka Govedič <luka.govedic@gmail.com>	2026-05-17 19:49:16 +00:00
TJianandGitHub	599e75f432	[ROCm] [Bugfix] Fix DeepSeek V4 Functionality and Accuracy (#42810 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-05-17 12:18:50 -04:00
Taneem IbrahimandGitHub	1c8e9c0399	Refactor: Pass num_labels explicitly to PoolerClassify instead of reading from global config (#42851 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>	2026-05-17 14:40:21 +00:00
zofia GitHub Kunshang Ji	0fa888465e	[XPU] fix weight scale shape (#42725 ) Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-05-17 16:55:10 +08:00
liuzhenweiandGitHub	ff712f6447	[MRV2][XPU] add Model Runner V2 log (#42710 ) Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>	2026-05-17 04:15:50 +00:00
Qi ZhouandGitHub	504a26ce2b	Support bf16 for mamba ssm cache (#41680 ) Signed-off-by: Qi Zhou <qizzzh@google.com>	2026-05-16 17:54:58 -07:00
weizhoublueandGitHub	a94189295b	Fix Weight loading for Qwen3.5-MTP and Qwen3-VL using runai_streamer (#42716 ) Signed-off-by: weizhoublue <weizhou.lan@daocloud.io>	2026-05-16 17:54:27 -07:00
Artem Perevedentsev GitHub Vadim Gimpelson	0867497368	[CI/Build] Bump flashinfer to v0.6.11.post2 (#41711 ) Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>	2026-05-16 14:55:12 -07:00
Zhewen Li GitHub Zhewen Li Claude Opus 4.7	36e74c9ea4	[KV Connector] Support disk offloading in MooncakeStoreConnector (#42689 ) Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-05-16 13:34:15 -07:00
Taneem IbrahimandGitHub	787bc0d031	Add unit tests for pooler activation functions (#42824 ) Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>	2026-05-16 14:58:16 -04:00
weizhoublueandGitHub	d1586e1a12	Fix: Propagate pinned model revisions into Ultravox secondary weight loading (#42830 )	2026-05-16 17:02:54 +00:00
Jiangyun ZhuandGitHub	8a56da3845	[Experimental] Breakable CUDA graph (#42304 ) Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>	2026-05-16 22:04:12 +08:00
Andreas KaratzasandGitHub	4db300e95f	[ROCm][CI] Removed problematic command override mechanism (#42807 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-05-16 17:35:05 +08:00
Zhewen Li GitHub Zhewen Li khluu	657b42b592	[Docker][KVConnector] Build mooncake-transfer-engine from source (#42114 ) Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: khluu <khluu000@gmail.com>	2026-05-16 00:26:25 -07:00
Jee Jee LiandGitHub	32b7177909	[LoRA][Bugfix] Dedup LoRA wrapping for modules referenced from multiple attribute paths (MoE gate) (#42757 ) Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>	2026-05-16 11:22:35 +08:00
DustHunter GitHub Qwen-Coder wang.yuqi	39c67d714e	fix: add API key authorization to /v2 endpoints (#42594 ) Signed-off-by: DustHunter <dusthunter@126.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>	2026-05-16 01:29:27 +00:00
Viktor Pus GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	87a2adcb43	[Misc] Add common random prefix option to structured-output serving benchmark (#41632 ) Signed-off-by: Viktor Pus <viktorpus@tenstorrent.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-05-16 00:44:48 +00:00
Michael GoinandGitHub	852f567444	[Bugfix] Respect explicit --kv-cache-dtype over checkpoint kv_cache_scheme (#42782 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-05-15 17:15:52 -07:00
Michael GoinandGitHub	b2a27b82d9	[Kernel][UX] Add `--linear-backend` arg for linear kernel selection (#39538 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2026-05-15 17:07:39 -07:00
Keyi LiandGitHub	d0921bafef	[Bugfix] Unwrap VLM wrappers for EPLB on Model Runner V2 (#42706 )	2026-05-16 07:20:33 +08:00
rasdani GitHub OpenAI Codex Roger Wang	1ccdf87507	[Bugfix] Fix layerwise reload alias-buffer corruption (#42481 ) Signed-off-by: rasdani <73563550+rasdani@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2026-05-15 15:20:53 -07:00
Rita BrugarolasandGitHub	bd9dbe6060	[ROCm][Bugfix] Fix fused_mla_dual_rms_norm for AITER API rename _fused_qk_rmsnorm (#42606 ) Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>	2026-05-15 14:50:03 -06:00
Michael Goin GitHub Claude	de2d76f352	[Build] Switch CUDA 12.9 wheel builds to PyTorch manylinux_2_28 base (#41668 ) Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-05-15 13:46:16 -07:00
Sergei Skvortsov GitHub Michael Goin	9a7a273dfe	Add HumanEval and GSM8K benchmarks to datasets (#42648 ) Signed-off-by: southfreebird <yvorott@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-05-15 13:01:21 -07:00
Lanze Liu GitHub Matthew Bonanni	b2c58ee942	[FlashAttn] Fix supports_kv_cache_dtype() accepting unhandled fp8 kv-cache dtype variants (#42685 ) Signed-off-by: Lanze Liu <lanzetech@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>	2026-05-15 15:34:59 -04:00
frida-anderssonandGitHub	4d67d3bde2	[ROCm] Restore fast top_k_per_row kernels for sparse MLA when topk_tokens=2048 (#42072 ) Signed-off-by: Frida Andersson <fanderss@amd.com>	2026-05-15 19:02:57 +00:00
Blake Ledden GitHub Claude Opus 4.6 Pavani Majety	06d020bb6e	[Bugfix] Fix SM121 (DGX Spark) exclusion from Marlin/CUTLASS FP8 paths (#35568 ) Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Pavani Majety <pmajety@nvidia.com>	2026-05-15 10:59:00 -07:00
chunxiaozhengandGitHub	f45c210885	[LMCacheMPConnector] Prioritize importing the lmcache_mp_connector from lmcache (#42596 ) Signed-off-by: idellzheng <idellzheng@tencent.com>	2026-05-15 17:46:31 +00:00
akii96andGitHub	be7a03ea65	[ROCm] Widen AITER fused AR RMSNorm 1-stage gate (#42409 ) Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>	2026-05-15 17:44:38 +00:00
Wentao Ye GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	6147c70224	[Model Runner v2] Support reload weights (sleep mode) (#42673 ) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-05-15 16:41:23 +00:00
PatchyTIS GitHub PatchouliTaisa	0162596603	[Model Runner V2] FP32 gumbel sampling. (#41775 ) Signed-off-by: PatchouliTaisa <patchychen@tencent.com> Co-authored-by: PatchouliTaisa <patchychen@tencent.com>	2026-05-15 09:20:08 -07:00
Markus Hartikainen GitHub clintg6 frida-andersson Cursor	46a95815d3	[ROCm][MLA] FP8 ASM prefill for AITER dense MLA backend on gfx950 (#42509 ) Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: clintg6 <clint.greene@amd.com> Co-authored-by: frida-andersson <frida.andersson@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-05-15 23:56:58 +08:00
BadrBasowidandGitHub	fb5bd03f51	[Perf] Set IR Op Priority Once at Worker Init (#42631 ) Signed-off-by: BadrBasowid <badr.basowid@gmail.com>	2026-05-15 15:56:13 +00:00
Mohammad Miadh AngkadandGitHub	ee58665aac	[Bugfix] Fix DeepGEMM context lens contiguity in MLA indexer (#42135 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>	2026-05-15 23:29:58 +08:00
Wentao YeandGitHub	491e8d8539	[Perf] Optimize MLA attention `_v_up_proj` bmm by removing additional copy (#42561 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-15 08:14:26 -07:00
Wentao YeandGitHub	af9616d845	[Model Runner V2] Fix kv_connector `pre_forward` order (#42676 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-15 08:13:59 -07:00
L.B.R.GitHub L.B.R.	d792d993c1	[ROCm] Widen OAI Triton MoE capability range to include gfx12 (RDNA4) (#37826 ) Signed-off-by: L.B.R. <lbr@mmonad.com> Co-authored-by: L.B.R. <lbr@mmonad.com>	2026-05-15 07:59:57 -07:00
Aaron HaoandGitHub	e0a45f1455	[Feat][RL] IPC weight sync optimizations: multigpu support and chunked packed tensors (#37476 ) Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: hao-aaron <ahao@anyscale.com>	2026-05-15 22:53:06 +08:00
Benjamin ChislettandGitHub	0fe7550254	[Bugfix] DFlash FP8 KV-Cache (#42692 ) Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>	2026-05-15 08:29:45 -06:00
Li, JiangandGitHub	95cfe102a5	[Bugfix] Ensure embeding model compilation on CPU (#42709 ) Signed-off-by: jiang1.li <jiang1.li@intel.com>	2026-05-15 18:58:19 +08:00
pmaybank GitHub TJian	1dc3fe08ea	gemma3 multi-gpu bug-fix (#42630 ) Signed-off-by: Philip Maybank <pmaybank@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>	2026-05-15 02:32:05 -07:00
Juan Pérez de Algaba GitHub Cyrus Leung	d26a28ab03	fix: propagate revision/code_revision pins to all artifact boundaries (#42616 ) Signed-off-by: jperezde <jperezde@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-05-15 02:31:54 -07:00
yewentao256	88d64448f4	Merge branch 'main' into wentao-optimize-model-runner-v2-sampler Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-03-16 14:57:34 +00:00
Wentao YeandGitHub	aef79c52cc	Merge branch 'main' into wentao-optimize-model-runner-v2-sampler	2026-03-05 10:47:03 -05:00
yewentao256	aeb03ccd3f	using dataclass and `plan` Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-03-04 21:01:38 +00:00
yewentao256	0d556d62ad	Merge branch 'main' into wentao-optimize-model-runner-v2-sampler Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-03-04 19:25:08 +00:00
Wentao YeandGitHub	1337c7bea5	Merge branch 'main' into wentao-optimize-model-runner-v2-sampler	2026-02-25 17:01:56 -05:00
yewentao256	666d1381db	optimize model-runner v2 sampler Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-02-24 18:27:26 +00:00